"""Atomic PostgreSQL activation for deterministic Booking Excel imports.""" from __future__ import annotations import hashlib import json import os from dataclasses import dataclass from datetime import datetime from typing import Any, Callable, Mapping, Optional from arr_ingestion.contracts import ArtifactRef from .excel import ( PROCESSOR_NAME, PROCESSOR_VERSION, RESULT_SCHEMA_VERSION, XLSX_MIME, ExcelDocument, build_excel_parse_result, canonical_json_bytes, rule_set_sha256, ) TARGET_DATABASE = "booking_test" DATABASE_ENV = "ARR_DATABASE_URL" ARTIFACT_STORAGE_PROVIDER = "oss" ARTIFACT_BUCKET_ALIAS = "arr-private" class BookingExcelRepositoryError(RuntimeError): def __init__(self, code: str, safe_message: str, status: int = 503): super().__init__(safe_message) self.code = code self.safe_message = safe_message self.status = status @dataclass(frozen=True) class DatabaseConfig: dsn: str @classmethod def from_environment(cls) -> "DatabaseConfig": dsn = os.environ.get(DATABASE_ENV, "").strip() if not dsn: raise BookingExcelRepositoryError( "BOOKING_EXCEL_DATABASE_CONFIG_MISSING", "Booking Excel 数据库配置缺失", ) return cls(dsn=dsn) @dataclass(frozen=True) class BookingSourceSummary: source_batch_id: int source_kind: str filename: Optional[str] source_rows: int worksheet_count: int distinct_group_codes: int room_quantity: int activated_at: Optional[datetime] disposition: str = "current" def to_dict(self) -> dict[str, object]: return { "source_batch_id": self.source_batch_id, "source_type": ( "excel" if self.source_kind == "booking_excel" else "historical" ), "filename": self.filename if self.source_kind == "booking_excel" else None, "source_rows": self.source_rows, "worksheet_count": self.worksheet_count, "distinct_group_codes": self.distinct_group_codes, "room_quantity": self.room_quantity, "activated_at": ( self.activated_at.isoformat() if self.activated_at is not None else None ), "disposition": self.disposition, } CURRENT_SOURCE_SQL = """ SELECT batch.id, batch.source_kind, CASE WHEN batch.source_kind = 'booking_excel' THEN artifact.original_filename ELSE NULL END AS uploaded_filename, count(source.id)::bigint AS source_rows, count(DISTINCT source.source_worksheet)::bigint AS worksheet_count, count(DISTINCT source.group_code_key)::bigint AS distinct_group_codes, COALESCE(sum(source.no_of_rooms), 0)::bigint AS room_quantity, active_source.activated_at FROM booking.current_source_batch AS active_source JOIN booking.source_batches AS batch ON batch.id = active_source.source_batch_id JOIN ingestion.artifacts AS artifact ON artifact.id = batch.source_artifact_id LEFT JOIN booking.source_rows AS source ON source.source_batch_id = batch.id WHERE active_source.singleton AND batch.batch_status = 'accepted' GROUP BY batch.id, batch.source_kind, artifact.original_filename, active_source.activated_at """.strip() BATCH_SUMMARY_SQL = """ SELECT batch.id, batch.source_kind, CASE WHEN batch.source_kind = 'booking_excel' THEN artifact.original_filename ELSE NULL END AS uploaded_filename, count(source.id)::bigint AS source_rows, count(DISTINCT source.source_worksheet)::bigint AS worksheet_count, count(DISTINCT source.group_code_key)::bigint AS distinct_group_codes, COALESCE(sum(source.no_of_rooms), 0)::bigint AS room_quantity, active_source.activated_at FROM booking.source_batches AS batch JOIN ingestion.artifacts AS artifact ON artifact.id = batch.source_artifact_id LEFT JOIN booking.source_rows AS source ON source.source_batch_id = batch.id LEFT JOIN booking.current_source_batch AS active_source ON active_source.source_batch_id = batch.id AND active_source.singleton WHERE batch.id = %s AND batch.batch_status = 'accepted' GROUP BY batch.id, batch.source_kind, artifact.original_filename, active_source.activated_at """.strip() def _default_connect(dsn: str) -> Any: try: import psycopg # type: ignore[import-not-found] except ImportError: raise BookingExcelRepositoryError( "BOOKING_EXCEL_DATABASE_DRIVER_UNAVAILABLE", "Booking Excel 数据库驱动不可用", ) from None try: return psycopg.connect(dsn, autocommit=False) except Exception: raise BookingExcelRepositoryError( "BOOKING_EXCEL_DATABASE_UNAVAILABLE", "Booking Excel 数据库暂不可用", ) from None class PostgresBookingExcelRepository: def __init__( self, config: DatabaseConfig, connect: Optional[Callable[[str], Any]] = None, ) -> None: self._config = config self._connect = connect or _default_connect def _open(self) -> Any: try: return self._connect(self._config.dsn) except BookingExcelRepositoryError: raise except Exception: raise BookingExcelRepositoryError( "BOOKING_EXCEL_DATABASE_UNAVAILABLE", "Booking Excel 数据库暂不可用", ) from None @staticmethod def _begin(cursor: Any, *, read_only: bool = False) -> None: mode = "REPEATABLE READ READ ONLY" if read_only else "SERIALIZABLE" cursor.execute(f"SET TRANSACTION ISOLATION LEVEL {mode}") cursor.execute("SET LOCAL lock_timeout = '10s'") cursor.execute("SET LOCAL statement_timeout = '120s'") cursor.execute( "SELECT current_database(), to_regclass('booking.current_source_batch')" ) row = cursor.fetchone() if not row or row[0] != TARGET_DATABASE or row[1] is None: raise BookingExcelRepositoryError( "BOOKING_EXCEL_DATABASE_NOT_READY", "Booking Excel 数据库迁移尚未就绪", ) @staticmethod def _summary_from_row( row: object, *, disposition: str = "current", ) -> Optional[BookingSourceSummary]: if row is None: return None values = tuple(row) # type: ignore[arg-type] return BookingSourceSummary( source_batch_id=int(values[0]), source_kind=str(values[1]), filename=str(values[2]) if values[2] else None, source_rows=int(values[3]), worksheet_count=int(values[4]), distinct_group_codes=int(values[5]), room_quantity=int(values[6]), activated_at=values[7] if isinstance(values[7], datetime) else None, disposition=disposition, ) def current_source(self) -> Optional[BookingSourceSummary]: connection = self._open() try: with connection.cursor() as cursor: self._begin(cursor, read_only=True) cursor.execute(CURRENT_SOURCE_SQL) return self._summary_from_row(cursor.fetchone()) except BookingExcelRepositoryError: raise except Exception: raise BookingExcelRepositoryError( "BOOKING_EXCEL_DATABASE_UNAVAILABLE", "当前 Excel 数据源暂时无法读取", ) from None finally: try: connection.rollback() finally: connection.close() @staticmethod def _find_existing_batch(cursor: Any, source_sha256: str) -> Optional[int]: cursor.execute( """ SELECT batch.id FROM booking.source_batches AS batch JOIN ingestion.artifacts AS artifact ON artifact.id = batch.source_artifact_id WHERE batch.source_kind = 'booking_excel' AND batch.batch_status = 'accepted' AND artifact.artifact_kind = 'booking_excel' AND artifact.sha256 = %s ORDER BY batch.id LIMIT 1 FOR UPDATE OF batch """, (source_sha256,), ) row = cursor.fetchone() return int(row[0]) if row is not None else None @staticmethod def _activate(cursor: Any, source_batch_id: int) -> bool: cursor.execute( """ SELECT source_batch_id FROM booking.current_source_batch WHERE singleton FOR UPDATE """ ) current = cursor.fetchone() already_active = current is not None and int(current[0]) == source_batch_id if not already_active: cursor.execute( """ INSERT INTO booking.current_source_batch ( singleton, source_batch_id, activated_at ) VALUES (true, %s, now()) ON CONFLICT (singleton) DO UPDATE SET source_batch_id = EXCLUDED.source_batch_id, activated_at = EXCLUDED.activated_at """, (source_batch_id,), ) return already_active @staticmethod def _batch_summary( cursor: Any, source_batch_id: int, disposition: str, ) -> BookingSourceSummary: cursor.execute(BATCH_SUMMARY_SQL, (source_batch_id,)) summary = PostgresBookingExcelRepository._summary_from_row( cursor.fetchone(), disposition=disposition, ) if summary is None: raise BookingExcelRepositoryError( "BOOKING_EXCEL_DATABASE_STATE_INVALID", "Booking Excel 数据源状态无效", ) return summary def activate_existing( self, source_sha256: str, ) -> Optional[BookingSourceSummary]: connection = self._open() try: with connection.cursor() as cursor: self._begin(cursor) cursor.execute( "SELECT pg_advisory_xact_lock(hashtextextended(%s, 0))", ("arr_booking_excel_import",), ) source_batch_id = self._find_existing_batch(cursor, source_sha256) if source_batch_id is None: connection.rollback() return None already_active = self._activate(cursor, source_batch_id) summary = self._batch_summary( cursor, source_batch_id, "already_active" if already_active else "reactivated", ) connection.commit() return summary except BookingExcelRepositoryError: connection.rollback() raise except Exception: connection.rollback() raise BookingExcelRepositoryError( "BOOKING_EXCEL_DATABASE_WRITE_FAILED", "Excel 已校验,但数据源未能完成入库", ) from None finally: connection.close() @staticmethod def _ensure_artifact(cursor: Any, reference: ArtifactRef) -> int: cursor.execute( """ SELECT id, artifact_kind, original_filename, sha256, byte_size, mime_type FROM ingestion.artifacts WHERE storage_provider = %s AND bucket_alias = %s AND object_key = %s AND object_version_id IS NULL FOR SHARE """, ( ARTIFACT_STORAGE_PROVIDER, ARTIFACT_BUCKET_ALIAS, reference.object_key, ), ) row = cursor.fetchone() if row is not None: if ( 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 str(row[5] or "") != reference.mime_type ): raise BookingExcelRepositoryError( "BOOKING_EXCEL_ARTIFACT_CONFLICT", "Excel 工件身份冲突", 409, ) return int(row[0]) cursor.execute( """ INSERT INTO ingestion.artifacts ( artifact_kind, storage_provider, bucket_alias, object_key, original_filename, sha256, byte_size, mime_type ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s) RETURNING id """, ( reference.file_kind, ARTIFACT_STORAGE_PROVIDER, ARTIFACT_BUCKET_ALIAS, reference.object_key, reference.original_filename, reference.sha256, reference.byte_size, reference.mime_type, ), ) return int(cursor.fetchone()[0]) def import_document( self, job_id: str, source: ArtifactRef, document: ExcelDocument, ) -> BookingSourceSummary: if ( source.role != "booking_source" or source.file_kind != "booking_excel" or source.mime_type != XLSX_MIME or source.sha256 != document.source_sha256 or source.byte_size != document.source_byte_size ): raise BookingExcelRepositoryError( "BOOKING_EXCEL_ARTIFACT_INVALID", "Excel 工件身份无效", 422, ) if document.pending_item_count: raise BookingExcelRepositoryError( "BOOKING_EXCEL_REVIEW_REQUIRED", "仍有房型需要人工确认或删除", 409, ) connection = self._open() try: with connection.cursor() as cursor: self._begin(cursor) cursor.execute( "SELECT pg_advisory_xact_lock(hashtextextended(%s, 0))", ("arr_booking_excel_import",), ) existing_batch_id = self._find_existing_batch( cursor, document.source_sha256, ) if existing_batch_id is not None: already_active = self._activate(cursor, existing_batch_id) summary = self._batch_summary( cursor, existing_batch_id, "already_active" if already_active else "reactivated", ) connection.commit() return summary artifact_id = self._ensure_artifact(cursor, source) delivery = { "source_sha256": document.source_sha256, "source_rows": len(document.rows), "worksheet_count": document.worksheet_count, "distinct_group_codes": document.distinct_group_code_count, "room_quantity": document.room_quantity, } delivery_bytes = canonical_json_bytes(delivery) delivery_sha256 = hashlib.sha256(delivery_bytes).hexdigest() cursor.execute( """ INSERT INTO ingestion.processing_runs ( run_key, pipeline_type, source_artifact_id, run_status, requested_processor_version, requested_rule_set_sha256, delivered_processor_version, delivered_rule_set_sha256, result_schema_version, delivery_sha256, delivery_json ) VALUES ( %s, 'booking_source_import', %s, 'validating', %s, %s, %s, %s, %s, %s, %s::jsonb ) RETURNING id """, ( job_id, artifact_id, PROCESSOR_VERSION, rule_set_sha256(), PROCESSOR_VERSION, rule_set_sha256(), RESULT_SCHEMA_VERSION, delivery_sha256, delivery_bytes.decode("utf-8"), ), ) processing_run_id = int(cursor.fetchone()[0]) cursor.execute( """ INSERT INTO booking.source_batches ( source_artifact_id, source_kind, source_format_version, batch_status, source_rows, accepted_rows, failed_rows ) VALUES ( %s, 'booking_excel', 'xlsx/1.0', 'processing', %s, 0, 0 ) RETURNING id """, (artifact_id, len(document.rows)), ) source_batch_id = int(cursor.fetchone()[0]) for row in document.rows: row_sha256 = row.sha256() result = build_excel_parse_result(row) result_bytes = canonical_json_bytes(result) result_sha256 = hashlib.sha256(result_bytes).hexdigest() cursor.execute( """ INSERT INTO booking.source_rows ( source_batch_id, source_worksheet, source_row_no, group_code_raw, type_of_room_raw, no_of_rooms, source_row_sha256 ) VALUES (%s, %s, %s, %s, %s, %s, %s) RETURNING id """, ( source_batch_id, row.worksheet, row.row_no, row.group_code_raw, row.type_of_room_raw, row.no_of_rooms, row_sha256, ), ) source_row_id = int(cursor.fetchone()[0]) cursor.execute( """ INSERT INTO booking.parse_versions ( source_row_id, version_no, parse_status, result_schema_version, processor_name, processor_version, rule_set_sha256, input_row_sha256, result_sha256, result_json, validated_at ) VALUES ( %s, 1, 'accepted', %s, %s, %s, %s, %s, %s, %s::jsonb, now() ) RETURNING id """, ( source_row_id, RESULT_SCHEMA_VERSION, PROCESSOR_NAME, PROCESSOR_VERSION, rule_set_sha256(), row_sha256, result_sha256, result_bytes.decode("utf-8"), ), ) parse_version_id = int(cursor.fetchone()[0]) for item in row.room_items: cursor.execute( """ INSERT INTO booking.room_items ( parse_version_id, item_no, room_type_raw, room_type_code, quantity, unit_price, currency_code, price_token_raw, source_fragment ) VALUES (%s, %s, %s, %s, %s, NULL, NULL, NULL, %s) """, ( parse_version_id, item.item_no, item.room_type_raw, item.room_type_code, item.quantity, item.source_fragment, ), ) cursor.execute( """ UPDATE booking.source_batches SET batch_status = 'accepted', accepted_rows = source_rows, failed_rows = 0, validated_at = now(), finished_at = now() WHERE id = %s """, (source_batch_id,), ) cursor.execute( """ INSERT INTO booking.current_row_parses ( source_row_id, parse_version_id ) SELECT source.id, parsed.id FROM booking.source_rows AS source JOIN booking.parse_versions AS parsed ON parsed.source_row_id = source.id AND parsed.version_no = 1 AND parsed.parse_status = 'accepted' WHERE source.source_batch_id = %s """, (source_batch_id,), ) self._activate(cursor, source_batch_id) cursor.execute( """ UPDATE ingestion.processing_runs SET run_status = 'accepted', updated_at = now(), validated_at = now(), finished_at = now() WHERE id = %s """, (processing_run_id,), ) cursor.execute( """ INSERT INTO ingestion.outbox_events ( event_key, aggregate_type, aggregate_id, event_type, payload ) VALUES ( %s, 'booking_source_batch', %s, 'booking.source_batch.accepted', %s::jsonb ) """, ( f"booking-source-batch:{source_batch_id}:accepted", source_batch_id, json.dumps( { "source_batch_id": source_batch_id, "source_rows": len(document.rows), "distinct_group_codes": document.distinct_group_code_count, }, separators=(",", ":"), ), ), ) summary = self._batch_summary( cursor, source_batch_id, "imported_and_activated", ) connection.commit() return summary except BookingExcelRepositoryError: connection.rollback() raise except Exception: connection.rollback() raise BookingExcelRepositoryError( "BOOKING_EXCEL_DATABASE_WRITE_FAILED", "Excel 已校验,但数据源未能完成入库", ) from None finally: connection.close()