feat: sync latest ARR implementation

This commit is contained in:
Wyndham ARR
2026-07-31 15:11:42 +08:00
parent d6f8a747fa
commit bf7939dd1a
185 changed files with 17527 additions and 2260 deletions

View File

@@ -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'