feat: sync latest ARR implementation
This commit is contained in:
@@ -216,19 +216,34 @@ class PostgresIngestionRepository(IngestionRepository):
|
||||
def _ensure_artifact(cursor: Any, reference: ArtifactRef) -> int:
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT id, byte_size, mime_type
|
||||
SELECT
|
||||
id,
|
||||
artifact_kind,
|
||||
original_filename,
|
||||
sha256,
|
||||
byte_size,
|
||||
mime_type
|
||||
FROM ingestion.artifacts
|
||||
WHERE artifact_kind = %s
|
||||
AND sha256 = %s
|
||||
WHERE storage_provider = %s
|
||||
AND bucket_alias = %s
|
||||
AND object_key = %s
|
||||
AND object_version_id IS NULL
|
||||
FOR SHARE
|
||||
""",
|
||||
(reference.file_kind, reference.sha256),
|
||||
(
|
||||
ARTIFACT_STORAGE_PROVIDER,
|
||||
ARTIFACT_BUCKET_ALIAS,
|
||||
reference.object_key,
|
||||
),
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
if row:
|
||||
if (
|
||||
int(row[1]) != reference.byte_size
|
||||
or (row[2] or "") != reference.mime_type
|
||||
str(row[1]) != reference.file_kind
|
||||
or str(row[2]) != reference.original_filename
|
||||
or str(row[3]) != reference.sha256
|
||||
or int(row[4]) != reference.byte_size
|
||||
or (row[5] or "") != reference.mime_type
|
||||
):
|
||||
raise IngestionError(
|
||||
"ARTIFACT_CONFLICT",
|
||||
@@ -278,6 +293,127 @@ class PostgresIngestionRepository(IngestionRepository):
|
||||
"processing job could not be stored",
|
||||
)
|
||||
|
||||
def mark_running(self, job_id: str, attempt_no: int) -> None:
|
||||
self._run_transaction(
|
||||
lambda cursor: self._mark_running(cursor, job_id, attempt_no),
|
||||
"processing job could not be started",
|
||||
)
|
||||
|
||||
def _mark_running(self, cursor: Any, job_id: str, attempt_no: int) -> None:
|
||||
run_id, run_status, attempt_id, attempt_status = self._lock_attempt(
|
||||
cursor, job_id, attempt_no
|
||||
)
|
||||
if run_status in {"accepted", "rejected", "failed", "cancelled"}:
|
||||
raise IngestionError("JOB_TERMINAL", "processing job is already terminal")
|
||||
if attempt_status in {"succeeded", "failed", "cancelled"}:
|
||||
raise IngestionError("JOB_TERMINAL", "processing attempt is already terminal")
|
||||
cursor.execute(
|
||||
"""
|
||||
UPDATE ingestion.processing_attempts
|
||||
SET attempt_status = 'running',
|
||||
started_at = COALESCE(started_at, now())
|
||||
WHERE id = %s
|
||||
""",
|
||||
(attempt_id,),
|
||||
)
|
||||
cursor.execute(
|
||||
"""
|
||||
UPDATE ingestion.processing_runs
|
||||
SET run_status = 'running',
|
||||
failure_code = NULL,
|
||||
failure_message = NULL,
|
||||
updated_at = now()
|
||||
WHERE id = %s
|
||||
""",
|
||||
(run_id,),
|
||||
)
|
||||
|
||||
def record_failure(self, job_id: str, attempt_no: int, failure_code: str) -> None:
|
||||
if not isinstance(failure_code, str) or re.fullmatch(
|
||||
r"[A-Z][A-Z0-9_]{0,63}", failure_code
|
||||
) is None:
|
||||
failure_code = "PROCESSING_FAILED"
|
||||
self._run_transaction(
|
||||
lambda cursor: self._record_runtime_failure(
|
||||
cursor, job_id, attempt_no, failure_code
|
||||
),
|
||||
"processing failure could not be stored",
|
||||
)
|
||||
|
||||
def _record_runtime_failure(
|
||||
self,
|
||||
cursor: Any,
|
||||
job_id: str,
|
||||
attempt_no: int,
|
||||
failure_code: str,
|
||||
) -> None:
|
||||
run_id, run_status, attempt_id, _attempt_status = self._lock_attempt(
|
||||
cursor, job_id, attempt_no
|
||||
)
|
||||
if run_status == "failed":
|
||||
return
|
||||
if run_status in {"accepted", "rejected", "cancelled"}:
|
||||
raise IngestionError("JOB_TERMINAL", "processing job is already terminal")
|
||||
cursor.execute(
|
||||
"""
|
||||
UPDATE ingestion.processing_attempts
|
||||
SET attempt_status = 'failed',
|
||||
failure_code = %s,
|
||||
failure_message = 'programmatic processing failed',
|
||||
finished_at = now()
|
||||
WHERE id = %s
|
||||
""",
|
||||
(failure_code, attempt_id),
|
||||
)
|
||||
cursor.execute(
|
||||
"""
|
||||
UPDATE ingestion.processing_runs
|
||||
SET run_status = 'failed',
|
||||
failure_code = %s,
|
||||
failure_message = 'programmatic processing failed',
|
||||
updated_at = now(),
|
||||
finished_at = now()
|
||||
WHERE id = %s
|
||||
""",
|
||||
(failure_code, run_id),
|
||||
)
|
||||
self._insert_outbox(
|
||||
cursor,
|
||||
f"processing-run:{run_id}:failed",
|
||||
"processing_run",
|
||||
run_id,
|
||||
"arr.processing_failed",
|
||||
{
|
||||
"job_id": job_id,
|
||||
"business_date": None,
|
||||
"daily_version_id": None,
|
||||
"failure_code": failure_code,
|
||||
},
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _lock_attempt(
|
||||
cursor: Any,
|
||||
job_id: str,
|
||||
attempt_no: int,
|
||||
) -> Tuple[int, str, int, str]:
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT run.id, run.run_status, attempt.id, attempt.attempt_status
|
||||
FROM ingestion.processing_runs AS run
|
||||
JOIN ingestion.processing_attempts AS attempt
|
||||
ON attempt.processing_run_id = run.id
|
||||
AND attempt.attempt_no = %s
|
||||
WHERE run.run_key = %s
|
||||
FOR UPDATE OF run, attempt
|
||||
""",
|
||||
(attempt_no, job_id),
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
if not row:
|
||||
raise IngestionError("JOB_NOT_FOUND", "processing job was not registered")
|
||||
return int(row[0]), str(row[1]), int(row[2]), str(row[3])
|
||||
|
||||
def _register_job(self, cursor: Any, registration: JobRegistration) -> None:
|
||||
source_artifact_id = self._ensure_artifact(cursor, registration.source)
|
||||
cursor.execute(
|
||||
@@ -287,6 +423,7 @@ class PostgresIngestionRepository(IngestionRepository):
|
||||
source_artifact_id,
|
||||
requested_processor_version,
|
||||
requested_rule_set_sha256,
|
||||
uploaded_filename,
|
||||
run_status
|
||||
FROM ingestion.processing_runs
|
||||
WHERE run_key = %s
|
||||
@@ -309,9 +446,10 @@ class PostgresIngestionRepository(IngestionRepository):
|
||||
source_artifact_id,
|
||||
run_status,
|
||||
requested_processor_version,
|
||||
requested_rule_set_sha256
|
||||
requested_rule_set_sha256,
|
||||
uploaded_filename
|
||||
)
|
||||
VALUES (%s, 'opera_daily', %s, 'queued', %s, %s)
|
||||
VALUES (%s, 'opera_daily', %s, 'queued', %s, %s, %s)
|
||||
RETURNING id
|
||||
""",
|
||||
(
|
||||
@@ -319,6 +457,7 @@ class PostgresIngestionRepository(IngestionRepository):
|
||||
source_artifact_id,
|
||||
registration.processor_version,
|
||||
registration.rule_set_sha256,
|
||||
registration.uploaded_filename,
|
||||
),
|
||||
)
|
||||
processing_run_id = int(cursor.fetchone()[0])
|
||||
@@ -327,6 +466,7 @@ class PostgresIngestionRepository(IngestionRepository):
|
||||
int(existing[1]) != source_artifact_id
|
||||
or existing[2] != registration.processor_version
|
||||
or str(existing[3]) != registration.rule_set_sha256
|
||||
or existing[4] != registration.uploaded_filename
|
||||
):
|
||||
raise IngestionError(
|
||||
"JOB_CONFLICT",
|
||||
@@ -350,7 +490,7 @@ class PostgresIngestionRepository(IngestionRepository):
|
||||
"processing attempt identity conflicts",
|
||||
)
|
||||
return
|
||||
if str(existing[4]) in {
|
||||
if str(existing[5]) in {
|
||||
"accepted",
|
||||
"rejected",
|
||||
"failed",
|
||||
@@ -1121,6 +1261,9 @@ class PostgresIngestionRepository(IngestionRepository):
|
||||
source.group_code_key,
|
||||
count(DISTINCT source.id)
|
||||
FROM booking.source_rows AS source
|
||||
JOIN booking.current_source_batch AS active_source
|
||||
ON active_source.source_batch_id = source.source_batch_id
|
||||
AND active_source.singleton
|
||||
JOIN booking.source_batches AS batch
|
||||
ON batch.id = source.source_batch_id
|
||||
AND batch.batch_status = 'accepted'
|
||||
|
||||
Reference in New Issue
Block a user