import re

with open('backend/src/db/api/job_logs.js', 'r') as f:
    content = f.read()

content = content.replace(
    'await job_logs.setVehicle( data.vehicle || null, {\n            transaction,\n        });',
    'await job_logs.setVehicle( data.vehicle || null, {\n            transaction,\n        });\n        \n        await job_logs.setWorkersCompClass( data.workersCompClass || null, {\n            transaction,\n        });'
)

content = content.replace(
    '''if (data.vehicle !== undefined) {\n            await job_logs.setVehicle(\n              \n              data.vehicle,\n              \n              { transaction }\n            );\n        }''',
    '''if (data.vehicle !== undefined) {\n            await job_logs.setVehicle(\n              \n              data.vehicle,\n              \n              { transaction }\n            );\n        }\n        \n        if (data.workersCompClass !== undefined) {\n            await job_logs.setWorkersCompClass(\n              \n              data.workersCompClass,\n              \n              { transaction }\n            );\n        }'''
)

content = content.replace(
    '''output.vehicle = await job_logs.getVehicle({\n            transaction\n        });''',
    '''output.vehicle = await job_logs.getVehicle({\n            transaction\n        });\n        \n        output.workersCompClass = await job_logs.getWorkersCompClass({\n            transaction\n        });'''
)

with open('backend/src/db/api/job_logs.js', 'w') as f:
    f.write(content)

print("Patched correctly!")