DM-54070: Add support for APDB record updates

.github/workflows/build.yaml

@@ -36,12 +36,17 @@ jobs:
           mamba install -y -q pip wheel
           pip install uv
 
+      - name: Install Postgres for testing
+        shell: bash -l {0}
+        run: |
+          mamba install -y -q postgresql
+
       - name: Install dependencies
         shell: bash -l {0}
         run: |
           uv pip install -r requirements.txt
+          uv pip install testing.postgresql
 
-      # We have two cores so we can speed up the testing with xdist
       - name: Install pytest packages
         shell: bash -l {0}
         run: |

.gitignore

@@ -27,3 +27,6 @@ pytest_session.txt
 
 # VS Code
 .vscode
+
+# Scratch directory
+.scratch

docker/Dockerfile.replication

@@ -3,11 +3,14 @@ FROM python:3.12-slim-bookworm
 ENV DEBIAN_FRONTEND=noninteractive
 
 # Update and install OS dependencies
-RUN apt-get -y update && \
-    apt-get -y upgrade && \
-    apt-get -y install --no-install-recommends git && \
-    apt-get clean && \
-    rm -rf /var/lib/apt/lists/*
+RUN apt-get update \
+ && apt-get install -y --no-install-recommends \
+      build-essential \
+      python3-dev \
+      pkg-config \
+      git \
+ && apt-get clean \
+ && rm -rf /var/lib/apt/lists/*
 
 # Install required python build dependencies
 RUN pip install --upgrade --no-cache-dir pip setuptools wheel uv

pyproject.toml

@@ -23,10 +23,12 @@ classifiers = [
 keywords = ["lsst"]
 dependencies = [
     "astropy",
+    "google-cloud-bigquery",
     "pyarrow",
     "pydantic >=2,<3",
     "pyyaml >= 5.1",
     "sqlalchemy",
+    "lsst-dax-ppdbx-gcp",
     "lsst-felis",
     "lsst-sdm-schemas",
     "lsst-utils",
@@ -43,9 +45,6 @@ test = [
     "pytest >= 3.2",
     "pytest-openfiles >= 0.5.0"
 ]
-gcp = [
-    "lsst-dax-ppdbx-gcp"
-]
 
 [tool.setuptools.packages.find]
 where = ["python"]
@@ -54,7 +53,7 @@ where = ["python"]
 zip-safe = true
 
 [tool.setuptools.package-data]
-"lsst.dax.ppdb" = ["py.typed"]
+"lsst.dax.ppdb" = ["py.typed", "config/schemas/*.yaml", "config/sql/*.sql"]
 
 [tool.setuptools.dynamic]
 version = { attr = "lsst_versions.get_lsst_version" }

python/lsst/dax/ppdb/__init__.py

@@ -19,7 +19,7 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
-from .config import *
+from .ppdb_config import *
 from .ppdb import *
 from .replicator import *
 from .version import *  # Generated by sconsUtils

python/lsst/dax/ppdb/_factory.py

@@ -26,8 +26,8 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from .config import PpdbConfig
     from .ppdb import Ppdb
+    from .ppdb_config import PpdbConfig
 
 
 def config_type_for_name(type_name: str) -> type[PpdbConfig]:

python/lsst/dax/ppdb/bigquery/__init__.py

@@ -20,5 +20,6 @@
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
 from .manifest import Manifest
+from .chunk_uploader import ChunkUploader
 from .ppdb_bigquery import PpdbBigQuery, PpdbBigQueryConfig
 from .ppdb_replica_chunk_extended import ChunkStatus, PpdbReplicaChunkExtended

python/lsst/dax/ppdb/bigquery/chunk_uploader.py

@@ -42,7 +42,7 @@
     ) from e
 
 
-from ..config import PpdbConfig
+from ..ppdb_config import PpdbConfig
 from .manifest import Manifest
 from .ppdb_bigquery import PpdbBigQuery, PpdbBigQueryConfig
 from .ppdb_replica_chunk_extended import ChunkStatus, PpdbReplicaChunkExtended
@@ -237,15 +237,27 @@ def _process_chunk(self, replica_chunk: PpdbReplicaChunkExtended) -> None:
         )
 
         # Make a list of local parquet files to upload.
-        parquet_files = list(chunk_dir.glob("*.parquet"))
+        upload_file_list = list(chunk_dir.glob("*.parquet"))
+
+        # Include the update records file if the manifest indicates it should
+        # exist
+        if manifest.includes_update_records:
+            update_records_file = chunk_dir / "update_records.json"
+            if not update_records_file.exists():
+                raise ChunkUploadError(
+                    chunk_id,
+                    f"Manifest indicates update records are included but file does not exist: "
+                    f"{update_records_file}",
+                )
+            upload_file_list.append(update_records_file)
 
         # Check if the chunk is expected to be empty.
         is_empty = manifest.is_empty_chunk()
 
-        if not parquet_files and not is_empty:
+        if not upload_file_list and not is_empty:
             # There is a mismatch between the manifest and the actual files.
             # Some processing error may have occurred when exporting.
-            raise ChunkUploadError(chunk_id, f"No parquet files found in {chunk_dir} for non-empty chunk")
+            raise ChunkUploadError(chunk_id, f"No files found to upload in {chunk_dir} for non-empty chunk")
 
         # Check that all expected parquet files from the manifest are present.
         for table_name, table_stats in manifest.table_data.items():
@@ -258,24 +270,21 @@ def _process_chunk(self, replica_chunk: PpdbReplicaChunkExtended) -> None:
                     )
 
         try:
-            # 1) Upload parquet files, which will happen only for non-empty
-            # chunks.
-            if parquet_files:
-                gcs_names = {path: posixpath.join(gcs_prefix, path.name) for path in parquet_files}
+            # 1) Upload the files to GCS for non-empty chunks
+            if upload_file_list:
+                gcs_names = {path: posixpath.join(gcs_prefix, path.name) for path in upload_file_list}
                 try:
-                    _LOG.info(
-                        "Uploading %d parquet files to GCS under prefix: %s", len(gcs_names), gcs_prefix
-                    )
+                    _LOG.info("Uploading %d files to GCS under prefix: %s", len(gcs_names), gcs_prefix)
                     with Timer(
                         "upload_files_time", _MON, tags={"prefix": str(gcs_prefix), "chunk_id": str(chunk_id)}
                     ) as timer:
                         self.storage.upload_files(gcs_names)
-                        total_bytes = sum(p.stat().st_size for p in parquet_files)
+                        total_bytes = sum(p.stat().st_size for p in upload_file_list)
                         timer.add_values(file_count=len(gcs_names), total_bytes=total_bytes)
                 except* UploadError as eg:
                     raise ChunkUploadError(chunk_id, f"{len(eg.exceptions)} upload(s) failed") from eg
 
-            # 2) Upload manifest, even for empty chunks.
+            # 2) Upload manifest, even for empty chunks
             try:
                 self.storage.upload_from_string(
                     posixpath.join(gcs_prefix, replica_chunk.manifest_name),
@@ -284,22 +293,29 @@ def _process_chunk(self, replica_chunk: PpdbReplicaChunkExtended) -> None:
             except UploadError as e:
                 raise ChunkUploadError(chunk_id, "Manifest upload failed") from e
 
-            # 3) Update DB status, but not for empty chunks.
+            # Next two steps are inapplicable to empty chunks.
             if not is_empty:
+                # 3) Update status and GCS URI in the database
+                gcs_prefix = posixpath.join(self.bucket_name, gcs_prefix)
+                updated_replica_chunk = replica_chunk.with_new_status(ChunkStatus.UPLOADED).with_new_gcs_uri(
+                    f"gs://{gcs_prefix}"
+                )
                 try:
-                    self._bq.store_chunk(replica_chunk.with_new_status(ChunkStatus.UPLOADED), True)
+                    self._bq.store_chunk(updated_replica_chunk, True)
+                    _LOG.info(
+                        "Updated replica chunk %d in database with status 'uploaded' and GCS URI: %s",
+                        chunk_id,
+                        gcs_prefix,
+                    )
                 except Exception as e:
-                    raise ChunkUploadError(
-                        chunk_id, "failed to update replica chunk status in database"
-                    ) from e
+                    raise ChunkUploadError(chunk_id, "Failed to update replica chunk in database") from e
 
-            # 4) Publish Pub/Sub staging message to trigger BigQuery load, but
-            # not for empty chunks. (Empty chunks cannot be staged.)
-            if not is_empty:
+                # 4) Publish Pub/Sub event to trigger staging of the chunk in
+                # BigQuery
                 try:
                     self._post_to_stage_chunk_topic(self.bucket_name, gcs_prefix, chunk_id)
                 except Exception as e:
-                    raise ChunkUploadError(chunk_id, "failed to publish staging message") from e
+                    raise ChunkUploadError(chunk_id, "Failed to publish staging message") from e
 
         except ChunkUploadError as err:
             try:

python/lsst/dax/ppdb/bigquery/manifest.py

@@ -79,6 +79,10 @@ class Manifest(BaseModel):
     """Name of the compression format used for artifacts (e.g., "gzip",
     "zstd", "snappy", etc.)."""
 
+    includes_update_records: bool = False
+    """Whether the exported data includes update records (e.g., in a separate
+    file) or not (`bool`)."""
+
     @property
     def filename(self) -> str:
         """Generate the filename for this manifest based on the replica chunk
@@ -118,12 +122,15 @@ def from_json_file(cls, file_path: Path) -> Manifest:
 
     def is_empty_chunk(self) -> bool:
         """Check if the manifest represents an empty replica chunk in which
-        all tables have zero rows.
+        all tables have zero rows and no update records are included.
 
         Returns
         -------
         bool
-            `True` if all tables have zero rows, indicating an empty chunk,
-            `False` otherwise.
+            `True` if all tables have zero rows and no update records are
+            included, indicating an empty chunk, `False` otherwise.
         """
-        return all(table.row_count == 0 for table in self.table_data.values())
+        return (
+            all(table.row_count == 0 for table in self.table_data.values())
+            and not self.includes_update_records
+        )