Iain Sproat
77 lines
2.2 KiB
Python
77 lines
2.2 KiB
Python
import json
|
|
|
|
from asyncpg import Connection, connect
|
|
|
|
from ifc_importer.config import settings
|
|
from ifc_importer.domain import FileimportJob, JobStatus
|
|
|
|
|
|
async def setup_connection() -> Connection:
|
|
connection = await connect(settings.fileimport_queue_postgres_url)
|
|
await connection.set_type_codec(
|
|
"jsonb",
|
|
encoder=json.dumps,
|
|
decoder=json.loads,
|
|
schema="pg_catalog",
|
|
)
|
|
return connection
|
|
|
|
|
|
async def get_next_job(connection: Connection) -> FileimportJob | None:
|
|
"""Get a fileimport job from the connection."""
|
|
|
|
job = await connection.fetchrow(
|
|
"""
|
|
WITH next_job AS (
|
|
UPDATE background_jobs
|
|
SET
|
|
"attempt" = "attempt" + 1,
|
|
"status" = $1,
|
|
"updatedAt" = NOW()
|
|
WHERE id = (
|
|
SELECT id FROM background_jobs
|
|
WHERE ( --queued job
|
|
payload ->> 'fileType' = 'ifc'
|
|
AND status = $2
|
|
AND "attempt" < "maxAttempt"
|
|
)
|
|
OR ( --timed job left on processing state
|
|
payload ->> 'fileType' = 'ifc'
|
|
AND status = $1
|
|
AND "updatedAt" < NOW() - ("timeoutMs" * interval '1 millisecond')
|
|
)
|
|
ORDER BY "createdAt"
|
|
FOR UPDATE SKIP LOCKED
|
|
LIMIT 1
|
|
)
|
|
RETURNING *
|
|
)
|
|
SELECT * FROM next_job;
|
|
""",
|
|
JobStatus.PROCESSING.value,
|
|
JobStatus.QUEUED.value,
|
|
)
|
|
if not job:
|
|
return None
|
|
return FileimportJob.model_validate(dict(job))
|
|
|
|
|
|
async def return_job_to_queued(connection: Connection, job_id: str) -> None:
|
|
print(f"returning job: {job_id} to queued")
|
|
return await set_job_status(connection, job_id, JobStatus.QUEUED)
|
|
|
|
|
|
async def set_job_status(
|
|
connection: Connection, job_id: str, job_status: JobStatus
|
|
) -> None:
|
|
print(f"updating job: {job_id}'s status to {job_status}")
|
|
_ = await connection.execute(
|
|
"""
|
|
UPDATE background_jobs
|
|
SET status = $1, "updatedAt" = NOW()
|
|
WHERE id = $2
|
|
""",
|
|
job_status.value,
|
|
job_id,
|
|
)
|