"""Read-only PostgreSQL snapshot provider for generated company reports.""" from __future__ import annotations import os import re import hashlib from dataclasses import dataclass from datetime import date from decimal import Decimal from pathlib import PurePosixPath from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Protocol, Sequence, Tuple from company_reports.contracts import ( BatchSnapshot, BookingRoomItem, CompanyReport, DailyVersionPin, ErrorCode, FinanceFact, PROCESSOR_VERSION, RESULT_SCHEMA_VERSION, rule_set_sha256, ) from company_reports.core import normalize_group_code, report_month_bounds SOURCE_FACTS_SQL = """ SELECT business_date, daily_version_id, id AS daily_record_id, channel_key, company_key, company_name, block_code, group_code_key, res_comment, room_category_label AS opera_room_category_label, arrival, departure, nights, total_price, booking_source_match_status FROM finance.v_active_daily_facts WHERE departure BETWEEN %s AND %s AND ( UPPER(BTRIM(channel_key)) = 'DY-AI-EASY-KB' OR UPPER(BTRIM(company_key)) = ANY(%s) ) ORDER BY daily_record_id """.strip() CURRENT_PARSE_SQL = """ SELECT group_code_key, max(parse_version_id) AS group_snapshot_id FROM booking.v_current_room_items WHERE group_code_key = ANY(%s) GROUP BY group_code_key ORDER BY group_code_key """.strip() ROOM_ITEMS_SQL = """ SELECT group_code_key, booking_room_type, quantity FROM booking.v_group_room_item_summary WHERE group_code_key = ANY(%s) ORDER BY group_code_key, booking_room_type """.strip() COMPANY_KEYS: Tuple[str, ...] = ("LIAN TAI", "QBD", "FENGRUN", "HANA TOUR") SHA256_RE = re.compile(r"^[0-9a-f]{64}$") class RepositoryError(RuntimeError): def __init__(self, code: str, safe_message: str): super().__init__(safe_message) self.code = code self.safe_message = safe_message @dataclass(frozen=True) class DatabaseConfig: dsn: str @classmethod def from_environment(cls) -> "DatabaseConfig": dsn = os.environ.get("COMPANY_REPORT_DATABASE_URL", "").strip() if not dsn: dsn = os.environ.get("ARR_DATABASE_URL", "").strip() if not dsn: raise RepositoryError( ErrorCode.REQUEST_INVALID, "COMPANY_REPORT_DATABASE_URL or ARR_DATABASE_URL is required", ) return cls(dsn=dsn) @dataclass(frozen=True) class ReservedReport: report_version_id: int version_no: int company: str @dataclass(frozen=True) class FileMetadata: file_kind: str original_filename: str storage_key: str sha256: str byte_size: int mime_type: str def validate(self) -> None: storage = PurePosixPath(self.storage_key) if ( not self.storage_key or storage.is_absolute() or ".." in storage.parts or self.original_filename != PurePosixPath(self.original_filename).name or not SHA256_RE.fullmatch(self.sha256) or self.byte_size < 0 ): raise RepositoryError( ErrorCode.PUBLISH_FAILED, "generated artifact metadata is invalid", ) class ReportRepository(Protocol): def load_snapshot(self, year: int, month: int, as_of_date: date) -> BatchSnapshot: ... def reserve_report(self, report: CompanyReport) -> ReservedReport: ... def activate_report( self, reservation: ReservedReport, artifact: FileMetadata, result_json: FileMetadata, ) -> None: ... def mark_failed(self, reservation: ReservedReport, code: str, safe_message: str) -> None: ... def _default_connect(dsn: str) -> Any: try: import psycopg # type: ignore[import-not-found] except ImportError: raise RepositoryError( ErrorCode.INTERNAL_ERROR, "PostgreSQL driver is unavailable; install requirements-company-reports.txt", ) from None try: return psycopg.connect(dsn, autocommit=True) except Exception: raise RepositoryError( ErrorCode.INTERNAL_ERROR, "database connection failed", ) from None def _rows(cursor: Any) -> List[Tuple[Any, ...]]: return list(cursor.fetchall()) class PostgresReportRepository: """Loads one immutable source snapshot without persisting report rows.""" 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 RepositoryError: raise except Exception: raise RepositoryError( ErrorCode.INTERNAL_ERROR, "database connection failed", ) from None def load_snapshot(self, year: int, month: int, as_of_date: date) -> BatchSnapshot: month_start, _ = report_month_bounds(year, month) connection = self._open() try: with connection.transaction(): with connection.cursor() as cursor: cursor.execute( "SET TRANSACTION ISOLATION LEVEL REPEATABLE READ READ ONLY" ) cursor.execute( SOURCE_FACTS_SQL, (month_start, as_of_date, list(COMPANY_KEYS)), ) fact_rows = _rows(cursor) facts = tuple( FinanceFact( business_date=row[0], daily_version_id=int(row[1]), daily_record_id=int(row[2]), channel_key=str(row[3] or ""), company_key=str(row[4] or ""), company_name=str(row[5] or ""), block_code=row[6], group_code_key=row[7], res_comment=row[8], room_category_label=row[9], arrival=row[10], departure=row[11], nights=row[12], total_price=( row[13] if row[13] is None or isinstance(row[13], Decimal) else Decimal(str(row[13])) ), booking_source_match_status=str(row[14] or "not_checked"), ) for row in fact_rows ) group_codes = sorted( { normalized for normalized in ( normalize_group_code(fact.res_comment) for fact in facts ) if normalized } ) parse_versions: Dict[str, int] = {} room_items: Tuple[BookingRoomItem, ...] = tuple() if group_codes: cursor.execute(CURRENT_PARSE_SQL, (group_codes,)) for group_code, parse_version_id in _rows(cursor): normalized = normalize_group_code(group_code) if normalized in parse_versions: raise RepositoryError( ErrorCode.BOOKING_PARSE_FAILED, "multiple current booking parses exist for one Group Code", ) parse_versions[normalized] = int(parse_version_id) cursor.execute(ROOM_ITEMS_SQL, (group_codes,)) summaries: Dict[str, List[Tuple[str, int]]] = {} for group_code, room_type, quantity in _rows(cursor): normalized = normalize_group_code(group_code) summaries.setdefault(normalized, []).append( (str(room_type or ""), int(quantity)) ) segments: Dict[ str, set[Tuple[date, date, int]], ] = {} for fact in facts: normalized = normalize_group_code(fact.group_code_key) if ( normalized and isinstance(fact.arrival, date) and isinstance(fact.departure, date) and isinstance(fact.nights, int) ): segments.setdefault(normalized, set()).add( (fact.arrival, fact.departure, fact.nights) ) expanded: List[BookingRoomItem] = [] for group_code in sorted(summaries): parse_version_id = parse_versions.get(group_code) if parse_version_id is None: continue for segment_no, segment in enumerate( sorted(segments.get(group_code, set())), 1, ): arrival, departure, nights = segment for item_no, item in enumerate( summaries[group_code], 1, ): room_type, quantity = item expanded.append( BookingRoomItem( group_code_key=group_code, parse_version_id=parse_version_id, segment_no=segment_no, arrival=arrival, departure=departure, nights=nights, item_no=item_no, room_type_raw=room_type, quantity=quantity, ) ) room_items = tuple(expanded) pins_by_date: Dict[date, int] = {} for fact in facts: existing = pins_by_date.get(fact.business_date) if existing is not None and existing != fact.daily_version_id: raise RepositoryError( ErrorCode.SOURCE_VERSION_MISSING, "the source snapshot contains conflicting daily versions", ) pins_by_date[fact.business_date] = fact.daily_version_id pins = tuple( DailyVersionPin(business_date=current_date, daily_version_id=version_id) for current_date, version_id in sorted(pins_by_date.items()) ) return BatchSnapshot( facts=facts, room_items=room_items, daily_versions=pins, group_parse_versions=dict(sorted(parse_versions.items())), ) except RepositoryError: raise except Exception: raise RepositoryError( ErrorCode.INTERNAL_ERROR, "database snapshot query failed", ) from None finally: connection.close() def reserve_report(self, report: CompanyReport) -> ReservedReport: if not report.valid: raise RepositoryError( ErrorCode.OUTPUT_VALIDATION_FAILED, "an invalid company report cannot be reserved", ) period_start, period_end = report_month_bounds( report.report_year, report.report_month ) identity = ( report.company, period_start.isoformat(), period_end.isoformat(), report.as_of_date.isoformat(), PROCESSOR_VERSION, rule_set_sha256(), RESULT_SCHEMA_VERSION, tuple( (pin.business_date.isoformat(), pin.daily_version_id) for pin in report.daily_versions ), tuple(sorted(report.booking_versions.items())), ) generation_id = int( hashlib.sha256(repr(identity).encode("utf-8")).hexdigest()[:12], 16, ) return ReservedReport( report_version_id=generation_id, version_no=generation_id, company=report.company, ) def activate_report( self, reservation: ReservedReport, artifact: FileMetadata, result_json: FileMetadata, ) -> None: if not reservation.company: raise RepositoryError( ErrorCode.PUBLISH_FAILED, "company report source reservation is invalid", ) artifact.validate() result_json.validate() def mark_failed( self, reservation: ReservedReport, code: str, safe_message: str, ) -> None: # Company report rows and versions are intentionally not persisted. _ = (reservation, code, safe_message)