From 0ae314319807977aa8804f72b0084bcaa8dfc74a Mon Sep 17 00:00:00 2001 From: Jeremy McCormick Date: Mon, 13 Jul 2026 18:26:43 -0500 Subject: [PATCH] Change update records to be staged like regular table data This changes update records to be handled more like regular table data in the replication and data ingestion. Instead of the promotion process working directly from the parquet files in cloud storage, it now reads from a staging table that is built similarly to the other tables such as DiaObject and DiaSource. --- .../lsst/dax/ppdb/bigquery/chunk_uploader.py | 29 ++-- .../lsst/dax/ppdb/bigquery/ppdb_bigquery.py | 2 +- .../ppdb/bigquery/schema/dataset_builder.py | 12 ++ .../dax/ppdb/bigquery/updates/__init__.py | 3 +- .../updates/expanded_update_record.py | 37 +++- ...tes_table.py => expanded_updates_table.py} | 112 ++++++------ .../updates/update_record_expander.py | 111 ------------ .../ppdb/bigquery/updates/update_records.py | 28 +-- .../ppdb/bigquery/updates/updates_manager.py | 159 ++++++++---------- python/lsst/dax/ppdb/tests/_updates.py | 4 +- tests/test_chunk_promoter.py | 18 +- tests/test_chunk_uploader.py | 4 +- tests/test_cli.py | 8 +- tests/test_dataset_builder.py | 12 +- ...nder.py => test_expanded_update_record.py} | 65 +++---- ...able.py => test_expanded_updates_table.py} | 95 ++++++----- tests/test_update_records.py | 9 +- tests/test_updates_manager.py | 18 +- tests/test_updates_merger.py | 57 ++++--- 19 files changed, 374 insertions(+), 409 deletions(-) rename python/lsst/dax/ppdb/bigquery/updates/{updates_table.py => expanded_updates_table.py} (64%) delete mode 100644 python/lsst/dax/ppdb/bigquery/updates/update_record_expander.py rename tests/{test_update_record_expander.py => test_expanded_update_record.py} (81%) rename tests/{test_updates_table.py => test_expanded_updates_table.py} (68%) diff --git a/python/lsst/dax/ppdb/bigquery/chunk_uploader.py b/python/lsst/dax/ppdb/bigquery/chunk_uploader.py index 6d05429b..7765a5f4 100644 --- a/python/lsst/dax/ppdb/bigquery/chunk_uploader.py +++ b/python/lsst/dax/ppdb/bigquery/chunk_uploader.py @@ -230,8 +230,9 @@ def _process_chunk(self, replica_chunk: PpdbReplicaChunkExtended) -> None: if not upload_file_list and not manifest.is_empty_chunk: raise ChunkUploadError(chunk_id, f"No files found to upload in {chunk_dir} for non-empty chunk") + gcs_uri = posixpath.join(self.config.bucket_name, gcs_prefix) try: - # 1) Upload the parquet files to GCS for non-empty chunks. + # 1) Upload the parquet files to GCS if there are any. if upload_file_list: gcs_names = { file_path: posixpath.join(gcs_prefix, file_path.name) for file_path in upload_file_list @@ -251,8 +252,7 @@ def _process_chunk(self, replica_chunk: PpdbReplicaChunkExtended) -> None: except* UploadError as eg: raise ChunkUploadError(chunk_id, f"{len(eg.exceptions)} upload(s) failed") from eg - # 2) Upload manifest, even for empty chunks. - gcs_uri = posixpath.join(self.config.bucket_name, gcs_prefix) + # 2) Upload the manifest JSON file, even for empty chunks. try: self._storage.upload_from_string( posixpath.join(gcs_prefix, manifest_path.name), @@ -261,13 +261,11 @@ def _process_chunk(self, replica_chunk: PpdbReplicaChunkExtended) -> None: except UploadError as e: raise ChunkUploadError(chunk_id, "Manifest upload failed") from e - # Next two steps are inapplicable to empty chunks. if not manifest.is_empty_chunk: - # 3) Update the status and GCS URI in the database, marking - # chunks with no table data as "staged" since they don't need - # to go through the staging process in Dataflow. - status = ChunkStatus.UPLOADED if manifest.has_table_data else ChunkStatus.STAGED - updated_replica_chunk = replica_chunk.with_new_status(status).with_new_gcs_uri( + # 3) Update the database record for the chunk, setting status + # to 'staged' and populating the GCS URI. + new_status = ChunkStatus.UPLOADED + updated_replica_chunk = replica_chunk.with_new_status(new_status).with_new_gcs_uri( f"gs://{gcs_uri}" ) try: @@ -275,19 +273,18 @@ def _process_chunk(self, replica_chunk: PpdbReplicaChunkExtended) -> None: _LOG.info( "Updated replica chunk %d in database with status '%s' and GCS URI: %s", chunk_id, - status, + new_status, gcs_uri, ) except Exception as e: raise ChunkUploadError(chunk_id, "Failed to update replica chunk in database") from e # 4) Publish Pub/Sub event to trigger the Dataflow staging - # job for chunks with table data (skipped for updates-only). - if manifest.has_table_data: - try: - self._post_to_stage_chunk_topic(gcs_uri, chunk_id) - except Exception as e: - raise ChunkUploadError(chunk_id, "Failed to publish staging message") from e + # for chunks with table data. + try: + self._post_to_stage_chunk_topic(gcs_uri, chunk_id) + except Exception as e: + raise ChunkUploadError(chunk_id, "Failed to publish staging message") from e except ChunkUploadError as err: try: diff --git a/python/lsst/dax/ppdb/bigquery/ppdb_bigquery.py b/python/lsst/dax/ppdb/bigquery/ppdb_bigquery.py index 4a840e9f..cb8f18c4 100644 --- a/python/lsst/dax/ppdb/bigquery/ppdb_bigquery.py +++ b/python/lsst/dax/ppdb/bigquery/ppdb_bigquery.py @@ -781,7 +781,7 @@ def _handle_updates( The directory containing the chunk's data. """ update_records = UpdateRecords( - records=list(apdb_update_records), + records=[(replica_chunk.id, record) for record in apdb_update_records], ) parquet_path = chunk_dir / UpdateRecords.PARQUET_FILE_NAME update_records.write_parquet_file(parquet_path) diff --git a/python/lsst/dax/ppdb/bigquery/schema/dataset_builder.py b/python/lsst/dax/ppdb/bigquery/schema/dataset_builder.py index fbf31296..dae0e2b9 100644 --- a/python/lsst/dax/ppdb/bigquery/schema/dataset_builder.py +++ b/python/lsst/dax/ppdb/bigquery/schema/dataset_builder.py @@ -217,6 +217,18 @@ def build_tables(self) -> list[bigquery.Table]: ) _update_schema_fields(table, apdb_replica_chunk_field) + # Add the table which will hold the raw update records. + updates_table = bigquery.Table( + self._config.fqn_for(self.dataset_type, "updates"), + schema=[ + bigquery.SchemaField("update_time_ns", "INT64", mode="REQUIRED"), + bigquery.SchemaField("update_order", "INT64", mode="REQUIRED"), + bigquery.SchemaField("json_payload", "STRING", mode="REQUIRED"), + bigquery.SchemaField("apdb_replica_chunk", "INT64", mode="REQUIRED"), + ], + ) + tables.append(updates_table) + return tables diff --git a/python/lsst/dax/ppdb/bigquery/updates/__init__.py b/python/lsst/dax/ppdb/bigquery/updates/__init__.py index a3ceb984..6a06b480 100644 --- a/python/lsst/dax/ppdb/bigquery/updates/__init__.py +++ b/python/lsst/dax/ppdb/bigquery/updates/__init__.py @@ -22,5 +22,4 @@ from .expanded_update_record import * from .updates_merger import * from .update_records import * -from .update_record_expander import * -from .updates_table import * +from .expanded_updates_table import * diff --git a/python/lsst/dax/ppdb/bigquery/updates/expanded_update_record.py b/python/lsst/dax/ppdb/bigquery/updates/expanded_update_record.py index 7c50ae80..ee99dedf 100644 --- a/python/lsst/dax/ppdb/bigquery/updates/expanded_update_record.py +++ b/python/lsst/dax/ppdb/bigquery/updates/expanded_update_record.py @@ -27,6 +27,8 @@ from pydantic import BaseModel, Field +from lsst.dax.apdb.apdbUpdateRecord import ApdbUpdateRecord + class ExpandedUpdateRecord(BaseModel): """A single normalized (expanded) update row. @@ -62,7 +64,7 @@ class ExpandedUpdateRecord(BaseModel): description=("New value for the field."), ) - replica_chunk_id: int = Field( + apdb_replica_chunk: int = Field( ..., ge=0, description=("Source replica chunk identifier associated with this update."), @@ -76,3 +78,36 @@ class ExpandedUpdateRecord(BaseModel): update_order: int = Field( description=("Ordering value within the replica chunk or update batch."), ) + + @classmethod + def from_update_record( + cls, update_record: ApdbUpdateRecord, apdb_replica_chunk: int + ) -> list[ExpandedUpdateRecord]: + """Expand a single APDB update record into field-level update rows. + + Parameters + ---------- + update_record + The APDB update record to expand. + apdb_replica_chunk + The replica chunk ID associated with this update record. + + Returns + ------- + `list` [ `ExpandedUpdateRecord` ] + One expanded record per field being updated. + """ + table_name = update_record.apdb_table.name + record_id_values = tuple(value for _, value in update_record.record_id()) + return [ + cls( + table_name=table_name, + record_id=record_id_values, + field_name=field_name, + field_value=field_value, + apdb_replica_chunk=apdb_replica_chunk, + update_order=update_record.update_order, + update_time_ns=update_record.update_time_ns, + ) + for field_name, field_value in update_record.record_payload() + ] diff --git a/python/lsst/dax/ppdb/bigquery/updates/updates_table.py b/python/lsst/dax/ppdb/bigquery/updates/expanded_updates_table.py similarity index 64% rename from python/lsst/dax/ppdb/bigquery/updates/updates_table.py rename to python/lsst/dax/ppdb/bigquery/updates/expanded_updates_table.py index 8a92a75c..d65cc299 100644 --- a/python/lsst/dax/ppdb/bigquery/updates/updates_table.py +++ b/python/lsst/dax/ppdb/bigquery/updates/expanded_updates_table.py @@ -21,7 +21,7 @@ from __future__ import annotations -__all__ = ["UpdatesTable"] +__all__ = ["ExpandedUpdatesTable"] from collections.abc import Iterable from typing import Any @@ -31,9 +31,13 @@ from .expanded_update_record import ExpandedUpdateRecord -class UpdatesTable: - """Manage the BigQuery table used for inserting expanded update records - which contain one update per row. +class ExpandedUpdatesTable: + """Manage the BigQuery tables holding expanded update records. + + This manages two related tables in the promotion dataset: the + ``expanded_updates`` table, which contains one row per updated field, and + the ``latest_only`` table, which is derived from it and contains only the + latest update for each field. Parameters ---------- @@ -44,35 +48,38 @@ class UpdatesTable: dataset_id BigQuery dataset ID. table_name - Name of the updates table. Defaults to ``"updates"``. + Name of the expanded updates table. Defaults to + ``"expanded_updates"``. latest_only_table_name - Name of the latest-only updates table. Defaults to - ``"updates_latest_only"``. + Name of the latest-only updates table. Defaults to ``"latest_only"``. """ - _DEFAULT_TABLE_NAME: str = "updates" - _DEFAULT_LATEST_ONLY_TABLE_NAME: str = "updates_latest_only" + EXPANDED_UPDATES_NAME: str = "expanded_updates" + LATEST_ONLY_NAME: str = "latest_only" def __init__( self, client: bigquery.Client, project_id: str, dataset_id: str, - table_name: str | None = None, - latest_only_table_name: str | None = None, + expanded_updates_name: str | None = None, + latest_only_name: str | None = None, ) -> None: self._client: bigquery.Client = client - table_name = table_name or self._DEFAULT_TABLE_NAME - latest_only_table_name = latest_only_table_name or self._DEFAULT_LATEST_ONLY_TABLE_NAME - self._table_fqn = f"{project_id}.{dataset_id}.{table_name}" - self._latest_only_table_fqn = f"{project_id}.{dataset_id}.{latest_only_table_name}" + expanded_updates_name = expanded_updates_name or self.EXPANDED_UPDATES_NAME + latest_only_name = latest_only_name or self.LATEST_ONLY_NAME + self._expanded_updates_fqn = f"{project_id}.{dataset_id}.{expanded_updates_name}" + self._latest_only_fqn = f"{project_id}.{dataset_id}.{latest_only_name}" @property - def latest_only_table_fqn(self) -> str: - """Fully-qualified BigQuery latest-only table name in the form - ``"project.dataset.table"`` (`str`, read-only). - """ - return self._latest_only_table_fqn + def expanded_updates_fqn(self) -> str: + """Fully-qualified expanded updates table name (`str`, read-only).""" + return self._expanded_updates_fqn + + @property + def latest_only_fqn(self) -> str: + """Fully-qualified latest-only table name (`str`, read-only).""" + return self._latest_only_fqn @staticmethod def _make_record_key(record_id: Iterable[int]) -> str: @@ -90,15 +97,13 @@ def _make_record_key(record_id: Iterable[int]) -> str: """ return "-".join(str(x) for x in record_id) - @property - def table_fqn(self) -> str: - """Fully-qualified BigQuery table name in the form - ``"project.dataset.table"`` (`str`, read-only). - """ - return self._table_fqn + def create(self, drop_if_exists: bool = False) -> bigquery.Table: + """Create the expanded updates table. - def create(self) -> bigquery.Table: - """Create the updates table. + Parameters + ---------- + drop_if_exists + If True, drop the table if it already exists before creating it. Returns ------- @@ -109,50 +114,29 @@ def create(self) -> bigquery.Table: ------ google.api_core.exceptions.Conflict Raised if the table already exists. - - Notes - ----- - Schema: - - - table_name: STRING (REQUIRED) - - record_id: ARRAY (REQUIRED) - - record_key: STRING (REQUIRED) - - field_name: STRING (REQUIRED) - - value_json: JSON (REQUIRED) - - replica_chunk_id: INT64 (REQUIRED) - - update_order: INT64 (REQUIRED) - - update_time_ns: INT64 (REQUIRED) """ + if drop_if_exists: + self._client.delete_table(self._expanded_updates_fqn, not_found_ok=True) schema: list[bigquery.SchemaField] = [ bigquery.SchemaField("table_name", "STRING", mode="REQUIRED"), bigquery.SchemaField("record_id", "INT64", mode="REPEATED"), bigquery.SchemaField("record_key", "STRING", mode="REQUIRED"), bigquery.SchemaField("field_name", "STRING", mode="REQUIRED"), bigquery.SchemaField("value_json", "JSON", mode="REQUIRED"), - bigquery.SchemaField("replica_chunk_id", "INT64", mode="REQUIRED"), + bigquery.SchemaField("apdb_replica_chunk", "INT64", mode="REQUIRED"), bigquery.SchemaField("update_order", "INT64", mode="REQUIRED"), bigquery.SchemaField("update_time_ns", "INT64", mode="REQUIRED"), ] - - table = bigquery.Table(self._table_fqn, schema=schema) + table = bigquery.Table(self.expanded_updates_fqn, schema=schema) return self._client.create_table(table) - def drop(self) -> None: - """Drop the table if it exists.""" - self._client.delete_table(self._table_fqn, not_found_ok=True) - - def recreate(self) -> None: - """Drop the table if it exists and then create it.""" - self.drop() - self.create() - def insert(self, records: Iterable[ExpandedUpdateRecord]) -> bigquery.LoadJob: - """Insert `ExpandedUpdateRecord` rows into the updates table. + """Insert `ExpandedUpdateRecord` rows into the expanded updates table. Parameters ---------- records - Iterable of update records to insert. + Iterable of expanded update records to insert. Returns ------- @@ -176,7 +160,7 @@ def insert(self, records: Iterable[ExpandedUpdateRecord]) -> bigquery.LoadJob: "record_key": self._make_record_key(r.record_id), "field_name": r.field_name, "value_json": r.field_value, - "replica_chunk_id": r.replica_chunk_id, + "apdb_replica_chunk": r.apdb_replica_chunk, "update_order": r.update_order, "update_time_ns": r.update_time_ns, } @@ -185,7 +169,7 @@ def insert(self, records: Iterable[ExpandedUpdateRecord]) -> bigquery.LoadJob: job = self._client.load_table_from_json( rows, - self._table_fqn, + self._expanded_updates_fqn, job_config=bigquery.LoadJobConfig( write_disposition=bigquery.WriteDisposition.WRITE_APPEND, ), @@ -199,18 +183,18 @@ def insert(self, records: Iterable[ExpandedUpdateRecord]) -> bigquery.LoadJob: def create_latest_only(self) -> None: """Select only the latest update for each unique - ``(table_name, record_id, field_name)`` combination and write them to a - new table. + ``(table_name, record_key, field_name)`` combination and write them to + the latest-only table. Notes ----- This keeps only the latest record with an update on an identical - ``(table_name, record_id, field_name)``, based on the descending - ordering of ``replica_chunk_id``, ``update_time_ns``, and + ``(table_name, record_key, field_name)``, based on the descending + ordering of ``apdb_replica_chunk``, ``update_time_ns``, and ``update_order``. """ query = f""" - CREATE OR REPLACE TABLE `{self._latest_only_table_fqn}` + CREATE OR REPLACE TABLE `{self._latest_only_fqn}` AS SELECT * EXCEPT(row_num) FROM ( @@ -218,11 +202,11 @@ def create_latest_only(self) -> None: ROW_NUMBER() OVER ( PARTITION BY table_name, record_key, field_name ORDER BY - replica_chunk_id DESC, + apdb_replica_chunk DESC, update_time_ns DESC, update_order DESC ) as row_num - FROM `{self._table_fqn}` + FROM `{self._expanded_updates_fqn}` ) WHERE row_num = 1 """ diff --git a/python/lsst/dax/ppdb/bigquery/updates/update_record_expander.py b/python/lsst/dax/ppdb/bigquery/updates/update_record_expander.py deleted file mode 100644 index f1988b6a..00000000 --- a/python/lsst/dax/ppdb/bigquery/updates/update_record_expander.py +++ /dev/null @@ -1,111 +0,0 @@ -# This file is part of dax_ppdb. -# -# Developed for the LSST Data Management System. -# This product includes software developed by the LSST Project -# (https://www.lsst.org). -# See the COPYRIGHT file at the top-level directory of this distribution -# for details of code ownership. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU General Public License -# along with this program. If not, see . - -from __future__ import annotations - -__all__ = ["UpdateRecordExpander"] - -import logging - -from lsst.dax.apdb.apdbUpdateRecord import ApdbUpdateRecord - -from .expanded_update_record import ExpandedUpdateRecord -from .update_records import UpdateRecords - -_LOG = logging.getLogger(__name__) - - -class UpdateRecordExpander: - """Expand APDB update records into individual field-level updates for - BigQuery. - """ - - @classmethod - def expand_single_record( - cls, update_record: ApdbUpdateRecord, replica_chunk_id: int - ) -> list[ExpandedUpdateRecord]: - """Expand a single APDB update record into ExpandedUpdateRecord - objects. - - Parameters - ---------- - update_record - A single APDB update record to expand. - replica_chunk_id - The replica chunk ID associated with this update record. - - Returns - ------- - `list` [ `ExpandedUpdateRecord` ] - List of ExpandedUpdateRecord objects, one per field being updated. - """ - # Get the target table from the update record. - table_name = update_record.apdb_table.name - - # Create an ExpandedUpdateRecord for each field being updated. - expanded_records = [] - record_id_values = tuple(value for _, value in update_record.record_id()) - for field_name, field_value in update_record.record_payload(): - expanded_record = ExpandedUpdateRecord( - table_name=table_name, - record_id=record_id_values, - field_name=field_name, - field_value=field_value, - replica_chunk_id=replica_chunk_id, - update_order=update_record.update_order, - update_time_ns=update_record.update_time_ns, - ) - expanded_records.append(expanded_record) - - return expanded_records - - @classmethod - def expand_updates( - cls, update_records: UpdateRecords, replica_chunk_id: int - ) -> list[ExpandedUpdateRecord]: - """Expand the APDB update records into a list of individual updates. - - Parameters - ---------- - update_records - The APDB update records to expand. - replica_chunk_id - The replica chunk ID associated with these update records. - - Returns - ------- - `list` [ `ExpandedUpdateRecord` ] - A list of individual updates derived from the input update records. - """ - expanded_updates = [] - - for update_record in update_records.records: - expanded_records = cls.expand_single_record(update_record, replica_chunk_id) - expanded_updates.extend(expanded_records) - - _LOG.info( - "Created %d expanded records from %d update records in chunk %d", - len(expanded_updates), - len(update_records.records), - replica_chunk_id, - ) - - return expanded_updates diff --git a/python/lsst/dax/ppdb/bigquery/updates/update_records.py b/python/lsst/dax/ppdb/bigquery/updates/update_records.py index 886e2689..87940f0b 100644 --- a/python/lsst/dax/ppdb/bigquery/updates/update_records.py +++ b/python/lsst/dax/ppdb/bigquery/updates/update_records.py @@ -39,10 +39,12 @@ class UpdateRecords: Parameters ---------- records - List of APDB update records. + List of ``(apdb_replica_chunk, record)`` pairs, associating each APDB + update record with the replica chunk it belongs to. The chunk id is + stored per row so that files for multiple chunks can be concatenated. """ - PARQUET_FILE_NAME: ClassVar[str] = "update_records.parquet" + PARQUET_FILE_NAME: ClassVar[str] = "updates.parquet" """Name of the Parquet file with the updates.""" _PARQUET_SCHEMA: ClassVar[pyarrow.Schema] = pyarrow.schema( @@ -50,18 +52,20 @@ class UpdateRecords: pyarrow.field("update_time_ns", pyarrow.int64()), pyarrow.field("update_order", pyarrow.int32()), pyarrow.field("json_payload", pyarrow.string()), + pyarrow.field("apdb_replica_chunk", pyarrow.int64()), ] ) - def __init__(self, records: list[ApdbUpdateRecord]) -> None: + def __init__(self, records: list[tuple[int, ApdbUpdateRecord]]) -> None: self.records = records def write_parquet_file(self, path: Path) -> None: """Write the update records to a Parquet file. - Each record is stored with ``update_time_ns``, ``update_order``, and - ``json_payload`` columns, where ``json_payload`` contains the - serialized record data. + Each record is stored with ``update_time_ns``, ``update_order``, + ``json_payload``, and ``apdb_replica_chunk`` columns, where + ``json_payload`` contains the serialized record data and + ``apdb_replica_chunk`` identifies the source replica chunk. Parameters ---------- @@ -71,16 +75,19 @@ def write_parquet_file(self, path: Path) -> None: update_times: list[int] = [] update_orders: list[int] = [] json_payloads: list[str] = [] - for record in self.records: + apdb_replica_chunks: list[int] = [] + for apdb_replica_chunk, record in self.records: update_times.append(record.update_time_ns) update_orders.append(record.update_order) json_payloads.append(record.to_json()) + apdb_replica_chunks.append(apdb_replica_chunk) table = pyarrow.table( { "update_time_ns": update_times, "update_order": update_orders, "json_payload": json_payloads, + "apdb_replica_chunk": apdb_replica_chunks, }, schema=self._PARQUET_SCHEMA, ) @@ -118,13 +125,14 @@ def from_parquet_bytes(cls, data: bytes) -> UpdateRecords: The deserialized update records. """ table = parquet.read_table(BytesIO(data), schema=cls._PARQUET_SCHEMA) - records: list[ApdbUpdateRecord] = [] - for update_time_ns, update_order, json_payload in zip( + records: list[tuple[int, ApdbUpdateRecord]] = [] + for update_time_ns, update_order, json_payload, apdb_replica_chunk in zip( table.column("update_time_ns").to_pylist(), table.column("update_order").to_pylist(), table.column("json_payload").to_pylist(), + table.column("apdb_replica_chunk").to_pylist(), strict=True, ): record = ApdbUpdateRecord.from_json(update_time_ns, update_order, json_payload) - records.append(record) + records.append((apdb_replica_chunk, record)) return cls(records=records) diff --git a/python/lsst/dax/ppdb/bigquery/updates/updates_manager.py b/python/lsst/dax/ppdb/bigquery/updates/updates_manager.py index 5ed747e9..3966f3df 100644 --- a/python/lsst/dax/ppdb/bigquery/updates/updates_manager.py +++ b/python/lsst/dax/ppdb/bigquery/updates/updates_manager.py @@ -24,23 +24,22 @@ __all__ = ["UpdatesManager", "UpdatesManagerError"] import logging -import posixpath -import urllib from collections.abc import Sequence -from google.cloud import bigquery, storage +from google.cloud import bigquery + +from lsst.dax.apdb.apdbUpdateRecord import ApdbUpdateRecord from ..ppdb_bigquery_config import DatasetType, PpdbBigQueryConfig from ..ppdb_replica_chunk_extended import PpdbReplicaChunkExtended -from .update_record_expander import UpdateRecordExpander -from .update_records import UpdateRecords +from .expanded_update_record import ExpandedUpdateRecord +from .expanded_updates_table import ExpandedUpdatesTable from .updates_merger import ( DiaForcedSourceUpdatesMerger, DiaObjectUpdatesMerger, DiaSourceUpdatesMerger, UpdatesMerger, ) -from .updates_table import UpdatesTable _LOG = logging.getLogger(__name__) @@ -51,9 +50,13 @@ class UpdatesManagerError(Exception): class UpdatesManager: """Class responsible for managing the process of applying updates to the - PPDB database, including expanding them into a generic format from JSON and - inserting into the updates table, selecting only the latest updates to a - new table, and, finally, merging the updates into target BigQuery tables. + PPDB database. + + Raw update records are staged into the ``updates`` table in the staging + dataset by the Dataflow staging job. This manager reads those raw records, + expands them into field-level rows in the ``expanded_updates`` table in the + promotion dataset, selects only the latest update for each field into the + ``latest_only`` table, and finally merges those into the target tables. Parameters ---------- @@ -74,9 +77,7 @@ def __init__( target_dataset_fqn: str | None = None, mergers: Sequence[UpdatesMerger] | None = None, ) -> None: - # Get some necessary setup information from the config. - project_id = config.project_id - bucket_name = config.bucket_name + self._updates_table_fqn = config.fqn_for(DatasetType.STAGING, "updates") # Set the merger instances for handling each target table, falling back # to a default set of mergers if none were provided. @@ -89,18 +90,16 @@ def __init__( DiaForcedSourceUpdatesMerger(), ) - # Setup the updates table interface. self._bq_client = bigquery.Client() - self._updates_table = UpdatesTable( + + # The expanded and latest-only tables are built in the promotion + # dataset. + self._expanded_updates_table = ExpandedUpdatesTable( self._bq_client, - project_id, + config.project_id, config.datasets.promotion, ) - # Setup the GCS client and bucket. - self._gcs_client = storage.Client() - self._bucket = self._gcs_client.bucket(bucket_name) - # Set the target dataset FQN for the mergers to use. if target_dataset_fqn is None: # By default, use the promotion dataset as the merge target. @@ -135,97 +134,87 @@ def apply_updates(self, replica_chunks: Sequence[PpdbReplicaChunkExtended]) -> N total_update_count = sum(chunk.update_count for chunk in update_chunks) _LOG.info("Processing %d update records from %d chunks", total_update_count, len(update_chunks)) - # Recreate the updates table. + # Read the raw update records from the staging table and expand them + # into the promotion expanded updates table. try: - self._updates_table.recreate() + self._build_expanded_updates_table(update_chunks) except Exception as e: - raise UpdatesManagerError("Failed to recreate updates table") from e - - # Build the table with the expanded update records. - try: - self._build_updates_table(update_chunks) - except Exception as e: - raise UpdatesManagerError("Failed to build updates table") from e + raise UpdatesManagerError("Failed to build expanded updates table") from e # Select only the latest update records into a new table. try: - self._updates_table.create_latest_only() + self._expanded_updates_table.create_latest_only() except Exception as e: raise UpdatesManagerError("Failed to create latest-only updates table") from e # Merge the latest-only updates into the target tables. try: - self._merge_updates(self._updates_table.latest_only_table_fqn) + self._merge_updates(self._expanded_updates_table.latest_only_fqn) except Exception as e: raise UpdatesManagerError("Failed to merge updates into target tables") from e - # TODO: This method will be removed by DM-55338, as populating the updates - # table will be handled by the staging process. - def _build_updates_table(self, chunks: Sequence[PpdbReplicaChunkExtended]) -> None: - """Build the updates table by expanding the update records from the - replica chunks and inserting them. + def _build_expanded_updates_table(self, chunks: Sequence[PpdbReplicaChunkExtended]) -> None: + """Read raw update records for the given chunks from the staging + updates table, expand them, and load them into the promotion expanded + updates table. Parameters ---------- chunks Replica chunks with update_count > 0. """ - for chunk in chunks: - # A null value for the GCS URI should not be possible under - # normal circumstances but check anyways before proceeding so that - # strange errors are not encountered later. - if chunk.gcs_uri is None: - raise UpdatesManagerError(f"Replica chunk {chunk.id} does not have a GCS URI") + chunk_ids = [chunk.id for chunk in chunks] + raw_records = self._read_staged_updates(chunk_ids) - try: - # Parse the GCS URI. - parsed_uri = urllib.parse.urlparse(chunk.gcs_uri) + expanded_records: list[ExpandedUpdateRecord] = [] + for apdb_replica_chunk, record in raw_records: + expanded_records.extend(ExpandedUpdateRecord.from_update_record(record, apdb_replica_chunk)) - # Create the GCS bucket. - bucket_name = parsed_uri.netloc - bucket = self._gcs_client.bucket(bucket_name) + _LOG.info( + "Expanded %d raw update records into %d field-level records", + len(raw_records), + len(expanded_records), + ) - # Get the object prefix from the URI path, stripping any - # leading slash. - chunk_prefix = parsed_uri.path.lstrip("/") + # Create the new expanded updates table, dropping any existing one. + self._expanded_updates_table.create(drop_if_exists=True) - if chunk_prefix == "": - raise ValueError(f"GCS URI '{chunk.gcs_uri}' does not contain an object prefix") + # Insert the expanded records into the new table. + self._expanded_updates_table.insert(expanded_records) - except Exception as e: - raise UpdatesManagerError( - f"Failed to parse GCS URI '{chunk.gcs_uri}' for replica chunk {chunk.id}" - ) from e + def _read_staged_updates(self, apdb_replica_chunks: Sequence[int]) -> list[tuple[int, ApdbUpdateRecord]]: + """Read staged update records for the given replica chunks. - # Download the parquet file containing the update records from the - # bucket. - try: - object_name = posixpath.join(chunk_prefix, UpdateRecords.PARQUET_FILE_NAME) - if not bucket.blob(object_name).exists(): - raise ValueError(f"GCS object '{object_name}' does not exist in bucket '{bucket_name}'") - blob = bucket.blob(object_name) - content = blob.download_as_bytes() - except Exception as e: - raise UpdatesManagerError( - f"Failed to download update records for replica chunk {chunk.id} " - f"from GCS URI '{chunk.gcs_uri}'" - ) from e + Parameters + ---------- + apdb_replica_chunks + Replica chunk IDs whose update records should be read. + + Returns + ------- + `list` [ `tuple` [ `int`, `ApdbUpdateRecord` ] ] + Pairs of ``(apdb_replica_chunk, record)`` for each row, with the + record reconstructed from its stored JSON payload. + """ + if not apdb_replica_chunks: + return [] - # Expand the update records into the appropriate format and insert - # them into the updates table. - try: - update_records = UpdateRecords.from_parquet_bytes(content) - if not update_records: - raise ValueError( - f"Empty file downloaded from GCS URI '{chunk.gcs_uri}' for replica chunk {chunk.id}" - ) - expanded_update_records = UpdateRecordExpander.expand_updates(update_records, chunk.id) - self._updates_table.insert(expanded_update_records) - except Exception as e: - raise UpdatesManagerError( - f"Failed to build updates table for replica chunk {chunk.id} " - f"from GCS URI '{chunk.gcs_uri}'" - ) from e + query = f""" + SELECT * FROM `{self._updates_table_fqn}` + WHERE apdb_replica_chunk IN UNNEST(@ids) + """ + job_config = bigquery.QueryJobConfig( + query_parameters=[bigquery.ArrayQueryParameter("ids", "INT64", list(apdb_replica_chunks))] + ) + rows = self._bq_client.query(query, job_config=job_config).result() + + records: list[tuple[int, ApdbUpdateRecord]] = [] + for row in rows: + record = ApdbUpdateRecord.from_json( + row["update_time_ns"], row["update_order"], row["json_payload"] + ) + records.append((row["apdb_replica_chunk"], record)) + return records def _merge_updates(self, target_table_fqn: str) -> None: """Merge the latest-only updates into the target tables using the @@ -253,7 +242,7 @@ def _merge_updates(self, target_table_fqn: str) -> None: except Exception as e: raise UpdatesManagerError(f"Failed to merge updates using {type(merger).__name__}") from e - def _log_job_stats(self, job: bigquery.job.QueryJob, target_table_fqn: str) -> None: + def _log_job_stats(self, job: bigquery.QueryJob, target_table_fqn: str) -> None: """Log relevant statistics for the merge operation from a BigQuery job. Parameters diff --git a/python/lsst/dax/ppdb/tests/_updates.py b/python/lsst/dax/ppdb/tests/_updates.py index f989b3ba..f369b38a 100644 --- a/python/lsst/dax/ppdb/tests/_updates.py +++ b/python/lsst/dax/ppdb/tests/_updates.py @@ -33,7 +33,7 @@ from ..bigquery.updates import UpdateRecords -def _create_test_update_records() -> UpdateRecords: +def _create_test_update_records(apdb_replica_chunk: int = 0) -> UpdateRecords: """Create test UpdateRecords with sample ApdbUpdateRecord instances.""" records: list[ApdbUpdateRecord] = [] @@ -149,5 +149,5 @@ def _create_test_update_records() -> UpdateRecords: ) return UpdateRecords( - records=records, + records=[(apdb_replica_chunk, record) for record in records], ) diff --git a/tests/test_chunk_promoter.py b/tests/test_chunk_promoter.py index 1cc3f7de..ade830e0 100644 --- a/tests/test_chunk_promoter.py +++ b/tests/test_chunk_promoter.py @@ -38,7 +38,7 @@ PpdbReplicaChunkExtended, ) from lsst.dax.ppdb.bigquery.sql_resource import SqlResource -from lsst.dax.ppdb.bigquery.updates import ExpandedUpdateRecord, UpdateRecordExpander, UpdateRecords +from lsst.dax.ppdb.bigquery.updates import ExpandedUpdateRecord, UpdateRecords from lsst.dax.ppdb.replicator import Replicator from lsst.dax.ppdb.tests import ( PostgresMixin, @@ -187,6 +187,16 @@ def _load_chunk_to_staging(self, chunk: PpdbReplicaChunkExtended) -> None: f"apdb_replica_chunk IS NULL" ).result() + def _load_updates_to_staging(self) -> None: + """Load each chunk's updates parquet into the staging ``updates`` + table, simulating the Dataflow staging of the updates file. + """ + updates_table_fqn = self.config.fqn_for(DatasetType.STAGING, "updates") + for chunk in self.ppdb.query_chunks(): + update_path = self.ppdb.config.chunk_dir(chunk.id) / UpdateRecords.PARQUET_FILE_NAME + if update_path.exists(): + self._load_parquet_to_table(update_path, updates_table_fqn) + def _query_table(self, table_fqn: str) -> list[dict]: """Query all rows from a BQ table.""" rows = list(self.bq_client.query(f"SELECT * FROM `{table_fqn}`").result()) @@ -207,7 +217,8 @@ def _get_expanded_updates(self) -> list[ExpandedUpdateRecord]: update_path = self.ppdb.config.chunk_dir(chunk.id) / UpdateRecords.PARQUET_FILE_NAME if update_path.exists(): update_records = UpdateRecords.from_parquet_file(update_path) - expanded.extend(UpdateRecordExpander.expand_updates(update_records, chunk.id)) + for apdb_replica_chunk, record in update_records.records: + expanded.extend(ExpandedUpdateRecord.from_update_record(record, apdb_replica_chunk)) return expanded def _verify_promoted_data( @@ -278,6 +289,9 @@ def test_promote_chunks(self) -> None: self._load_chunk_to_staging(chunk) self.ppdb.update_chunks([chunk.with_new_status(ChunkStatus.STAGED)], fields={"status"}) + # Stage the raw update records into the staging updates table. + self._load_updates_to_staging() + # Snapshot the staged rows per table before promotion. staging_rows: dict[str, list[dict]] = {} for table_name in self.TABLE_NAMES: diff --git a/tests/test_chunk_uploader.py b/tests/test_chunk_uploader.py index b0f01718..bf45f185 100644 --- a/tests/test_chunk_uploader.py +++ b/tests/test_chunk_uploader.py @@ -82,12 +82,12 @@ def test_chunk_uploader(self) -> None: ).run() # Retrieve the update records file. - blobs = list(self.bucket.list_blobs(match_glob="**/update_records.parquet")) + blobs = list(self.bucket.list_blobs(match_glob=f"**/{UpdateRecords.PARQUET_FILE_NAME}")) update_records_files = [b.name for b in blobs] self.assertEqual( len(update_records_files), 1, - f"Expected exactly one update_records.parquet file in GCS, found " + f"Expected exactly one {UpdateRecords.PARQUET_FILE_NAME} file in GCS, found " f"{len(update_records_files)}: {update_records_files}", ) diff --git a/tests/test_cli.py b/tests/test_cli.py index b5bcc613..ecbe9615 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -77,9 +77,11 @@ def test_create_datasets(self) -> None: dataset_fqn = self.config.fqn_for(dataset_type) self.client.get_dataset(dataset_fqn) tables = list(self.client.list_tables(dataset_fqn)) - if dataset_type != DatasetType.PROMOTION: - # The staging, internal, and public datasets should each have - # three tables created by default. + if dataset_type == DatasetType.STAGING: + # Staging has the three DIA tables plus the raw updates table. + self.assertEqual(len(tables), 4) + elif dataset_type != DatasetType.PROMOTION: + # The internal and public datasets each have three tables. self.assertEqual(len(tables), 3) else: # Promotion dataset should have no tables created by default. diff --git a/tests/test_dataset_builder.py b/tests/test_dataset_builder.py index 1a005fb9..1d210ac0 100644 --- a/tests/test_dataset_builder.py +++ b/tests/test_dataset_builder.py @@ -293,8 +293,15 @@ def test_adds_replica_chunk(self) -> None: tables = builder.build_tables() - self.assertEqual({table.table_id for table in tables}, self.EXPECTED_TABLES) + self.assertEqual( + {table.table_id for table in tables}, + self.EXPECTED_TABLES | {"updates"}, + ) for table in tables: + if table.table_id == "updates": + # The raw updates table is not one of the DIA tables that the + # builder adds apdb_replica_chunk to. + continue chunk_field = next(field for field in table.schema if field.name == "apdb_replica_chunk") self.assertEqual(chunk_field.field_type, "INT64") self.assertEqual(chunk_field.mode, "REQUIRED") @@ -719,7 +726,8 @@ def test_build_datasets_creates_expected_bigquery_objects(self) -> None: self.assertEqual(len(internal_tables), 3) self.assertEqual(len(public_tables), 3) - self.assertEqual(len(staging_tables), 3) + # Staging has the three DIA tables plus the raw updates table. + self.assertEqual(len(staging_tables), 4) # Public dataset should have a properly structured DiaObject table. public_dia_object = self.client.get_table( diff --git a/tests/test_update_record_expander.py b/tests/test_expanded_update_record.py similarity index 81% rename from tests/test_update_record_expander.py rename to tests/test_expanded_update_record.py index b8f0668c..4a53bb81 100644 --- a/tests/test_update_record_expander.py +++ b/tests/test_expanded_update_record.py @@ -31,12 +31,12 @@ ApdbWithdrawDiaForcedSourceRecord, ApdbWithdrawDiaSourceRecord, ) -from lsst.dax.ppdb.bigquery.updates import ExpandedUpdateRecord, UpdateRecordExpander, UpdateRecords +from lsst.dax.ppdb.bigquery.updates import ExpandedUpdateRecord from lsst.dax.ppdb.tests._updates import _create_test_update_records -class UpdateRecordExpanderTestCase(unittest.TestCase): - """Test UpdateRecordExpander functionality.""" +class ExpandedUpdateRecordTestCase(unittest.TestCase): + """Test ExpandedUpdateRecord.from_update_record functionality.""" def setUp(self) -> None: """Set up test fixtures.""" @@ -48,11 +48,9 @@ def setUp(self) -> None: self.replica_chunk_id = 12345 def test_reassign_diasource_to_diaobject(self) -> None: - """Test expand_single_record with + """Test from_update_record with ApdbReassignDiaSourceToDiaObjectRecord. """ - from lsst.dax.ppdb.bigquery.updates import ExpandedUpdateRecord, UpdateRecordExpander - record = ApdbReassignDiaSourceToDiaObjectRecord( update_time_ns=self.update_time_ns, update_order=0, @@ -63,7 +61,7 @@ def test_reassign_diasource_to_diaobject(self) -> None: midpointMjdTai=60000.0, ) - expanded = UpdateRecordExpander.expand_single_record(record, self.replica_chunk_id) + expanded = ExpandedUpdateRecord.from_update_record(record, self.replica_chunk_id) # Should expand to 1 record (diaObjectId) self.assertEqual(len(expanded), 1) @@ -74,14 +72,12 @@ def test_reassign_diasource_to_diaobject(self) -> None: self.assertEqual(expanded_record.record_id, (100001,)) self.assertEqual(expanded_record.field_name, "diaObjectId") self.assertEqual(expanded_record.field_value, 300001) - self.assertEqual(expanded_record.replica_chunk_id, self.replica_chunk_id) + self.assertEqual(expanded_record.apdb_replica_chunk, self.replica_chunk_id) self.assertEqual(expanded_record.update_order, 0) self.assertEqual(expanded_record.update_time_ns, self.update_time_ns) def test_reassign_diasource_to_ssobject(self) -> None: - """Test expand_single_record with - ApdbReassignDiaSourceToSSObjectRecord. - """ + """Test from_update_record with ApdbReassignDiaSourceToSSObjectRecord.""" record = ApdbReassignDiaSourceToSSObjectRecord( update_time_ns=self.update_time_ns, update_order=0, @@ -93,7 +89,7 @@ def test_reassign_diasource_to_ssobject(self) -> None: midpointMjdTai=60000.0, ) - expanded = UpdateRecordExpander.expand_single_record(record, self.replica_chunk_id) + expanded = ExpandedUpdateRecord.from_update_record(record, self.replica_chunk_id) # Should expand to 2 records (ssObjectId and ssObjectReassocTimeMjdTai) self.assertEqual(len(expanded), 2) @@ -105,7 +101,7 @@ def test_reassign_diasource_to_ssobject(self) -> None: self.assertEqual(first_record.record_id, (100001,)) self.assertEqual(first_record.field_name, "ssObjectId") self.assertEqual(first_record.field_value, 2001) - self.assertEqual(first_record.replica_chunk_id, self.replica_chunk_id) + self.assertEqual(first_record.apdb_replica_chunk, self.replica_chunk_id) self.assertEqual(first_record.update_order, 0) self.assertEqual(first_record.update_time_ns, self.update_time_ns) @@ -117,7 +113,7 @@ def test_reassign_diasource_to_ssobject(self) -> None: self.assertEqual(second_record.field_value, float(self.update_time.tai.mjd)) def test_withdraw_diasource(self) -> None: - """Test expand_single_record with ApdbWithdrawDiaSourceRecord.""" + """Test from_update_record with ApdbWithdrawDiaSourceRecord.""" record = ApdbWithdrawDiaSourceRecord( update_time_ns=self.update_time_ns, update_order=2, @@ -128,7 +124,7 @@ def test_withdraw_diasource(self) -> None: midpointMjdTai=60000.0, ) - expanded = UpdateRecordExpander.expand_single_record(record, self.replica_chunk_id) + expanded = ExpandedUpdateRecord.from_update_record(record, self.replica_chunk_id) # Should expand to 1 record (timeWithdrawnMjdTai) self.assertEqual(len(expanded), 1) @@ -140,7 +136,7 @@ def test_withdraw_diasource(self) -> None: self.assertEqual(expanded_record.field_value, self.update_time.tai.mjd) def test_update_n_dia_sources(self) -> None: - """Test expand_single_record with ApdbUpdateNDiaSourcesRecord.""" + """Test from_update_record with ApdbUpdateNDiaSourcesRecord.""" record = ApdbUpdateNDiaSourcesRecord( update_time_ns=self.update_time_ns, update_order=5, @@ -150,7 +146,7 @@ def test_update_n_dia_sources(self) -> None: dec=-30.0, ) - expanded = UpdateRecordExpander.expand_single_record(record, self.replica_chunk_id) + expanded = ExpandedUpdateRecord.from_update_record(record, self.replica_chunk_id) # Should expand to 1 record (nDiaSources) self.assertEqual(len(expanded), 1) @@ -162,7 +158,7 @@ def test_update_n_dia_sources(self) -> None: self.assertEqual(expanded_record.field_value, 10) def test_close_diaobject_validity(self) -> None: - """Test expand_single_record with ApdbCloseDiaObjectValidityRecord.""" + """Test from_update_record with ApdbCloseDiaObjectValidityRecord.""" record = ApdbCloseDiaObjectValidityRecord( update_time_ns=self.update_time_ns, update_order=4, @@ -173,7 +169,7 @@ def test_close_diaobject_validity(self) -> None: dec=-30.0, ) - expanded = UpdateRecordExpander.expand_single_record(record, self.replica_chunk_id) + expanded = ExpandedUpdateRecord.from_update_record(record, self.replica_chunk_id) # Should expand to 2 records (validityEndMjdTai and nDiaSources) self.assertEqual(len(expanded), 2) @@ -194,7 +190,7 @@ def test_close_diaobject_validity(self) -> None: self.assertEqual(second_record.field_value, 5) def test_withdraw_diaforcedsource(self) -> None: - """Test expand_single_record with ApdbWithdrawDiaForcedSourceRecord.""" + """Test from_update_record with ApdbWithdrawDiaForcedSourceRecord.""" record = ApdbWithdrawDiaForcedSourceRecord( update_time_ns=self.update_time_ns, update_order=2, @@ -207,30 +203,32 @@ def test_withdraw_diaforcedsource(self) -> None: midpointMjdTai=60000.0, ) - expanded = UpdateRecordExpander.expand_single_record(record, self.replica_chunk_id) + expanded = ExpandedUpdateRecord.from_update_record(record, self.replica_chunk_id) # Should expand to 1 record (timeWithdrawnMjdTai) self.assertEqual(len(expanded), 1) expanded_record = expanded[0] self.assertEqual(expanded_record.table_name, "DiaForcedSource") - # The record ID should be a list of the composite key components - # [diaObjectId, visit, detector] for BigQuery compatibility. + # The record ID should be the composite key components + # (diaObjectId, visit, detector) for BigQuery compatibility. expected_record_id = (200001, 12345, 42) self.assertEqual(expanded_record.record_id, expected_record_id) self.assertEqual(expanded_record.field_name, "timeWithdrawnMjdTai") self.assertEqual(expanded_record.field_value, self.update_time.tai.mjd) def test_update_records_all(self) -> None: - """Test the full expand_updates method with multiple record types.""" - update_records = _create_test_update_records() - expanded = UpdateRecordExpander.expand_updates(update_records, self.replica_chunk_id) + """Test expanding a full set of update records of multiple types.""" + update_records = _create_test_update_records(self.replica_chunk_id) + expanded: list[ExpandedUpdateRecord] = [] + for apdb_replica_chunk, record in update_records.records: + expanded.extend(ExpandedUpdateRecord.from_update_record(record, apdb_replica_chunk)) self.assertEqual(len(expanded), 10) - # Verify that all expanded records have correct the replica_chunk_id. + # Verify that all expanded records have the correct apdb_replica_chunk. for record in expanded: - self.assertEqual(record.replica_chunk_id, self.replica_chunk_id) + self.assertEqual(record.apdb_replica_chunk, self.replica_chunk_id) test_update_time_ns = 1640995200000000000 test_mjd_tai = 59580.0 @@ -277,11 +275,14 @@ def test_update_records_all(self) -> None: self.assertEqual(expanded[i].update_time_ns, update_time_ns) def test_empty_records(self) -> None: - """Test expand_updates with empty records list.""" - empty_update_records = UpdateRecords(records=[]) - - expanded = UpdateRecordExpander.expand_updates(empty_update_records, self.replica_chunk_id) + """Test that expanding an empty set of records yields nothing.""" + update_records = _create_test_update_records(self.replica_chunk_id) + expanded: list[ExpandedUpdateRecord] = [] + for apdb_replica_chunk, record in []: + expanded.extend(ExpandedUpdateRecord.from_update_record(record, apdb_replica_chunk)) self.assertEqual(len(expanded), 0) + # The helper itself should produce records to expand. + self.assertTrue(update_records.records) if __name__ == "__main__": diff --git a/tests/test_updates_table.py b/tests/test_expanded_updates_table.py similarity index 68% rename from tests/test_updates_table.py rename to tests/test_expanded_updates_table.py index c6b4c6ee..f2dbe69e 100644 --- a/tests/test_updates_table.py +++ b/tests/test_expanded_updates_table.py @@ -24,7 +24,7 @@ from google.cloud import bigquery from lsst.dax.ppdb.bigquery import DatasetType -from lsst.dax.ppdb.bigquery.updates import UpdateRecordExpander, UpdatesTable +from lsst.dax.ppdb.bigquery.updates import ExpandedUpdateRecord, ExpandedUpdatesTable from lsst.dax.ppdb.tests import ( create_datasets, drop_datasets, @@ -35,10 +35,10 @@ @unittest.skipIf(not have_valid_google_credentials(), "Missing valid Google credentials") -class TestUpdatesTable(unittest.TestCase): - """Test UpdatesTable functionality.""" +class TestExpandedUpdatesTable(unittest.TestCase): + """Test ExpandedUpdatesTable functionality.""" - dataset_types = (DatasetType.STAGING,) + dataset_types = (DatasetType.PROMOTION,) def setUp(self) -> None: """Set up test fixtures.""" @@ -46,7 +46,7 @@ def setUp(self) -> None: self.client = bigquery.Client() # Create PPDB BigQuery config for test. - self.config = make_bigquery_config(test_name="test_updates_table") + self.config = make_bigquery_config(test_name="test_expanded_updates_table") # Create the BigQuery dataset for the tests. create_datasets(self.config, self.dataset_types) @@ -54,23 +54,34 @@ def setUp(self) -> None: # Add cleanup for datasets after test. self.addCleanup(drop_datasets, self.config, self.dataset_types) - # Set the FQN of the updates table. - self.table_fqn = self.config.fqn_for(DatasetType.STAGING, "updates") + # Set the FQN of the expanded updates table. + self.table_fqn = self.config.fqn_for(DatasetType.PROMOTION, "expanded_updates") - # Create the UpdatesTable instance. - self.updates_table = UpdatesTable(self.client, self.config.project_id, self.config.datasets.staging) + # Create the ExpandedUpdatesTable instance. + self.expanded_updates_table = ExpandedUpdatesTable( + self.client, self.config.project_id, self.config.datasets.promotion + ) + + @staticmethod + def _expand(apdb_replica_chunk: int) -> list[ExpandedUpdateRecord]: + """Expand the test update records for the given chunk.""" + update_records = _create_test_update_records(apdb_replica_chunk) + expanded: list[ExpandedUpdateRecord] = [] + for chunk, record in update_records.records: + expanded.extend(ExpandedUpdateRecord.from_update_record(record, chunk)) + return expanded def test_table_fqn_property(self) -> None: """Test the table_fqn property.""" - self.assertEqual(self.updates_table.table_fqn, self.table_fqn) + self.assertEqual(self.expanded_updates_table.expanded_updates_fqn, self.table_fqn) def test_create_table(self) -> None: - """Test creation of the updates table.""" - table = self.updates_table.create() + """Test creation of the expanded updates table.""" + table = self.expanded_updates_table.create() # Verify table was created successfully. - self.assertEqual(table.table_id, "updates") - self.assertEqual(table.dataset_id, self.config.datasets.staging) + self.assertEqual(table.table_id, "expanded_updates") + self.assertEqual(table.dataset_id, self.config.datasets.promotion) # Verify schema is correct. expected_fields = { @@ -79,37 +90,35 @@ def test_create_table(self) -> None: "record_key": ("STRING", "REQUIRED"), "field_name": ("STRING", "REQUIRED"), "value_json": ("JSON", "REQUIRED"), - "replica_chunk_id": ("INTEGER", "REQUIRED"), + "apdb_replica_chunk": ("INTEGER", "REQUIRED"), "update_order": ("INTEGER", "REQUIRED"), "update_time_ns": ("INTEGER", "REQUIRED"), } - actual_fields = {field.name: (field.field_type, field.mode) for field in table.schema} self.assertEqual(actual_fields, expected_fields) def test_create_table_already_exists(self) -> None: """Test creating a table that already exists raises an error.""" # Create table first time - should succeed. - self.updates_table.create() + self.expanded_updates_table.create() # Try to create again - should raise Conflict. with self.assertRaises(Exception) as cm: - self.updates_table.create() + self.expanded_updates_table.create() # Check that it's a conflict-type error. self.assertIn("already exists", str(cm.exception).lower()) def test_insert_records(self) -> None: """Test insertion of expanded records into the table.""" - # Create the table first - self.updates_table.create() + # Create the table first. + self.expanded_updates_table.create() # Get test update records and expand them. - update_records = _create_test_update_records() - expanded_records = UpdateRecordExpander.expand_updates(update_records, 12345) + expanded_records = self._expand(12345) # Insert the records. - job = self.updates_table.insert(expanded_records) + job = self.expanded_updates_table.insert(expanded_records) # Verify the job completed successfully. self.assertIsNone(job.errors) @@ -119,46 +128,45 @@ def test_insert_records(self) -> None: result = list(self.client.query(query).result()) record_count = result[0].count - # Should have 10 total expanded records based on the test data - # (1 + 2 + 1 + 1 + 2 + 1 from original records + 2 duplicates) + # Should have 10 total expanded records based on the test data. self.assertEqual(record_count, 10) - # Verify some specific data was inserted correctly + # Verify some specific data was inserted correctly. query = f""" - SELECT table_name, record_id, field_name, replica_chunk_id + SELECT table_name, record_id, field_name, apdb_replica_chunk FROM `{self.table_fqn}` ORDER BY table_name, field_name, record_key """ results = list(self.client.query(query).result()) - # Verify record counts per table + # Verify record counts per table. tables = [r.table_name for r in results] self.assertEqual(tables.count("DiaSource"), 5) self.assertEqual(tables.count("DiaObject"), 4) self.assertEqual(tables.count("DiaForcedSource"), 1) - # All records should have the same replica_chunk_id + # All records should have the same apdb_replica_chunk. for row in results: - self.assertEqual(row.replica_chunk_id, 12345) + self.assertEqual(row.apdb_replica_chunk, 12345) - # Verify the single DiaForcedSource record + # Verify the single DiaForcedSource record. forced = [r for r in results if r.table_name == "DiaForcedSource"] self.assertEqual(len(forced), 1) self.assertEqual(list(forced[0].record_id), [200001, 12345, 42]) self.assertEqual(forced[0].field_name, "timeWithdrawnMjdTai") - # Verify the DiaSource records + # Verify the DiaSource records. reassign = [r for r in results if r.table_name == "DiaSource" and r.field_name == "diaObjectId"] self.assertEqual(len(reassign), 2) self.assertTrue(all(list(r.record_id) == [100001] for r in reassign)) def test_insert_empty_records(self) -> None: """Test insertion of empty record list.""" - # Create the table first - self.updates_table.create() + # Create the table first. + self.expanded_updates_table.create() # Insert empty list. - job = self.updates_table.insert([]) + job = self.expanded_updates_table.insert([]) # Verify the job completed successfully. self.assertIsNone(job.errors) @@ -172,34 +180,33 @@ def test_insert_empty_records(self) -> None: def test_latest_updates_only(self) -> None: """Test functionality for getting only the latest updates.""" # Create the source table. - self.updates_table.create() + self.expanded_updates_table.create() # Get test records and expand them. - update_records = _create_test_update_records() - expanded_records = UpdateRecordExpander.expand_updates(update_records, 0) + expanded_records = self._expand(0) - # Insert all of the records into the updates table. - self.updates_table.insert(expanded_records) + # Insert all of the records into the expanded updates table. + self.expanded_updates_table.insert(expanded_records) # Count the original records. query = f"SELECT COUNT(*) as count FROM `{self.table_fqn}`" original_count = list(self.client.query(query).result())[0].count # Create table with only the latest updates. - self.updates_table.create_latest_only() + self.expanded_updates_table.create_latest_only() # Count the new number of records. - query = f"SELECT COUNT(*) as count FROM `{self.updates_table.latest_only_table_fqn}`" + query = f"SELECT COUNT(*) as count FROM `{self.expanded_updates_table.latest_only_fqn}`" latest_only_count = list(self.client.query(query).result())[0].count # There should be fewer records now. self.assertLess(latest_only_count, original_count) # Verify specific record has the later update value. - record_key = UpdatesTable._make_record_key([100001]) + record_key = ExpandedUpdatesTable._make_record_key([100001]) query = f""" SELECT value_json - FROM `{self.updates_table.latest_only_table_fqn}` + FROM `{self.expanded_updates_table.latest_only_fqn}` WHERE record_key = '{record_key}' AND field_name = 'diaObjectId' """ result = list(self.client.query(query).result()) diff --git a/tests/test_update_records.py b/tests/test_update_records.py index 8686daf9..8e0adfc6 100644 --- a/tests/test_update_records.py +++ b/tests/test_update_records.py @@ -86,14 +86,15 @@ def test_parquet_serialization(self) -> None: self.assertEqual(len(update_records.records), 3, "Unexpected number of update records deserialized") - for record in update_records.records: + for apdb_replica_chunk, record in update_records.records: + self.assertEqual(apdb_replica_chunk, 1614600000) self.assertIsInstance( record, apdbUpdateRecord.ApdbUpdateRecord, "Deserialized record is not an instance of ApdbUpdateRecord", ) - update_record = update_records.records[0] + update_record = update_records.records[0][1] self.assertIsInstance( update_record, apdbUpdateRecord.ApdbReassignDiaSourceToSSObjectRecord, @@ -141,7 +142,7 @@ def test_parquet_serialization(self) -> None: "Unexpected dec in deserialized ApdbReassignDiaSourceToSSObjectRecord, should not be 0.0", ) - update_record = update_records.records[1] + update_record = update_records.records[1][1] self.assertIsInstance( update_record, apdbUpdateRecord.ApdbCloseDiaObjectValidityRecord, @@ -182,7 +183,7 @@ def test_parquet_serialization(self) -> None: "Unexpected nDiaSources in deserialized ApdbCloseDiaObjectValidityRecord, expected None", ) - update_record = update_records.records[2] + update_record = update_records.records[2][1] self.assertIsInstance( update_record, apdbUpdateRecord.ApdbWithdrawDiaForcedSourceRecord, diff --git a/tests/test_updates_manager.py b/tests/test_updates_manager.py index 5063f474..988ab9a0 100644 --- a/tests/test_updates_manager.py +++ b/tests/test_updates_manager.py @@ -39,6 +39,7 @@ PpdbBigQueryConfig, ) from lsst.dax.ppdb.bigquery.chunk_uploader import ChunkUploader +from lsst.dax.ppdb.bigquery.updates import UpdateRecords from lsst.dax.ppdb.bigquery.updates.updates_manager import UpdatesManager from lsst.dax.ppdb.tests._bigquery import ( PostgresMixin, @@ -208,7 +209,7 @@ def rows(self) -> Collection[tuple]: objects=DummyApdbTableData(), sources=DummyApdbTableData(), forced_sources=DummyApdbTableData(), - update_records=update_records.records, + update_records=[record for _, record in update_records.records], update=True, ) @@ -221,6 +222,21 @@ def rows(self) -> Collection[tuple]: exit_on_error=True, ).run() + # Stage the raw update records into the staging updates table, + # simulating the Dataflow staging of the updates file. + staging_client = bigquery.Client() + updates_parquet = self.ppdb.config.chunk_dir(test_replica_chunk_id) / UpdateRecords.PARQUET_FILE_NAME + with open(updates_parquet, "rb") as f: + load_job = staging_client.load_table_from_file( + f, + self.config.fqn_for(DatasetType.STAGING, "updates"), + job_config=bigquery.LoadJobConfig( + source_format=bigquery.SourceFormat.PARQUET, + write_disposition=bigquery.WriteDisposition.WRITE_APPEND, + ), + ) + load_job.result() + # Apply the updates to the target tables using the UpdatesManager. updates_manager = UpdatesManager(self.ppdb.config) replica_chunks = self.ppdb.query_chunks( diff --git a/tests/test_updates_merger.py b/tests/test_updates_merger.py index 0991d4c4..e212ed3f 100644 --- a/tests/test_updates_merger.py +++ b/tests/test_updates_merger.py @@ -28,8 +28,8 @@ DiaForcedSourceUpdatesMerger, DiaObjectUpdatesMerger, DiaSourceUpdatesMerger, - UpdateRecordExpander, - UpdatesTable, + ExpandedUpdateRecord, + ExpandedUpdatesTable, ) from lsst.dax.ppdb.tests import ( create_datasets, @@ -57,6 +57,22 @@ def setUp(self): create_datasets(self.config, self.dataset_types) + def _build_latest_only(self) -> str: + """Build the expanded and latest-only updates tables from the test + update records and return the latest-only table FQN. + """ + expanded_table = ExpandedUpdatesTable( + self.client, self.config.project_id, self.config.datasets.staging + ) + expanded_table.create() + update_records = _create_test_update_records() + expanded = [] + for apdb_replica_chunk, record in update_records.records: + expanded.extend(ExpandedUpdateRecord.from_update_record(record, apdb_replica_chunk)) + expanded_table.insert(expanded) + expanded_table.create_latest_only() + return expanded_table.latest_only_fqn + def _create_target_table(self): schema = [ bigquery.SchemaField("diaObjectId", "INTEGER", mode="REQUIRED"), @@ -81,19 +97,14 @@ def _create_target_table(self): def test_merge_diaobject(self): self._create_target_table() - updates_table = UpdatesTable(self.client, self.config.project_id, self.config.datasets.staging) - updates_table.create() - update_records = _create_test_update_records() - expanded = UpdateRecordExpander.expand_updates(update_records, 0) - updates_table.insert(expanded) - updates_table.create_latest_only() + latest_only_fqn = self._build_latest_only() table_fqn = self.config.fqn_for(DatasetType.INTERNAL, "DiaObject") query = f"SELECT * FROM `{table_fqn}` ORDER BY diaObjectId" before = {r.diaObjectId: r for r in self.client.query(query).result()} merger = DiaObjectUpdatesMerger() merger.merge( client=self.client, - updates_table_fqn=updates_table.latest_only_table_fqn, + updates_table_fqn=latest_only_fqn, target_dataset_fqn=self.config.fqn_for(DatasetType.INTERNAL), ) after = {r.diaObjectId: r for r in self.client.query(query).result()} @@ -152,19 +163,14 @@ def test_merge_diasource(self): ) job.result() - updates_table = UpdatesTable(self.client, self.config.project_id, self.config.datasets.staging) - updates_table.create() - update_records = _create_test_update_records() - expanded = UpdateRecordExpander.expand_updates(update_records, 0) - updates_table.insert(expanded) - updates_table.create_latest_only() + latest_only_fqn = self._build_latest_only() query = f"SELECT * FROM `{table_fqn}` ORDER BY diaSourceId" before = {r.diaSourceId: r for r in self.client.query(query).result()} merger = DiaSourceUpdatesMerger() merger.merge( client=self.client, - updates_table_fqn=updates_table.latest_only_table_fqn, + updates_table_fqn=latest_only_fqn, target_dataset_fqn=self.config.fqn_for(DatasetType.INTERNAL), ) after = {r.diaSourceId: r for r in self.client.query(query).result()} @@ -198,19 +204,14 @@ def test_merge_diaforcedsource(self): ) job.result() - updates_table = UpdatesTable(self.client, self.config.project_id, self.config.datasets.staging) - updates_table.create() - update_records = _create_test_update_records() - expanded = UpdateRecordExpander.expand_updates(update_records, 0) - updates_table.insert(expanded) - updates_table.create_latest_only() + latest_only_fqn = self._build_latest_only() query = f"SELECT * FROM `{table_fqn}` ORDER BY diaObjectId, visit, detector" before = {(r.diaObjectId, r.visit, r.detector): r for r in self.client.query(query).result()} merger = DiaForcedSourceUpdatesMerger() merger.merge( client=self.client, - updates_table_fqn=updates_table.latest_only_table_fqn, + updates_table_fqn=latest_only_fqn, target_dataset_fqn=self.config.fqn_for(DatasetType.INTERNAL), ) after = {(r.diaObjectId, r.visit, r.detector): r for r in self.client.query(query).result()} @@ -223,15 +224,17 @@ def test_merge_diaforcedsource(self): def test_merge_no_updates(self): self._create_target_table() - updates_table = UpdatesTable(self.client, self.config.project_id, self.config.datasets.staging) - updates_table.create() - updates_table.create_latest_only() + expanded_table = ExpandedUpdatesTable( + self.client, self.config.project_id, self.config.datasets.staging + ) + expanded_table.create() + expanded_table.create_latest_only() table_fqn = self.config.fqn_for(DatasetType.INTERNAL, "DiaObject") before = {r.diaObjectId: r for r in self.client.query(f"SELECT * FROM `{table_fqn}`").result()} merger = DiaObjectUpdatesMerger() merger.merge( client=self.client, - updates_table_fqn=updates_table.latest_only_table_fqn, + updates_table_fqn=expanded_table.latest_only_fqn, target_dataset_fqn=self.config.fqn_for(DatasetType.INTERNAL), ) after = {r.diaObjectId: r for r in self.client.query(f"SELECT * FROM `{table_fqn}`").result()}