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()}