StoreFactory refactor (#191)

* Add storage.py and update storage config references

* Remove final_store implement use of factory

* Refactor template_utils.write_metadata to take a StoreFactory

* Remove outdated comment

Author: mosegontar

src/reformatters/__main__.py

@@ -9,10 +9,8 @@
 import reformatters.noaa.gefs.forecast_35_day.cli as noaa_gefs_forecast_35_day
 from reformatters.common import deploy
 from reformatters.common.config import Config
-from reformatters.common.dynamical_dataset import (
-    DynamicalDataset,
-    DynamicalDatasetStorageConfig,
-)
+from reformatters.common.dynamical_dataset import DynamicalDataset
+from reformatters.common.storage import DatasetFormat, StorageConfig
 from reformatters.contrib.noaa.ndvi_cdr.analysis import (
     NoaaNdviCdrAnalysisDataset,
 )
@@ -24,21 +22,23 @@
 )
 
 
-class SourceCoopDatasetStorageConfig(DynamicalDatasetStorageConfig):
+class SourceCoopDatasetStorageConfig(StorageConfig):
     """Configuration for the storage of a SourceCoop dataset."""
 
     base_path: str = "s3://us-west-2.opendata.source.coop/dynamical"
     k8s_secret_names: Sequence[str] = ["source-coop-key"]
+    format: DatasetFormat = DatasetFormat.ZARR3
 
 
-class UpstreamGriddedZarrsDatasetStorageConfig(DynamicalDatasetStorageConfig):
+class UpstreamGriddedZarrsDatasetStorageConfig(StorageConfig):
     """Configuration for storage in the Upstream gridded zarrs bucket."""
 
     # This bucket is actually an R2 bucket.
     # The R2 endpoint URL is stored within our k8s secret and will be set
     # when it's imported into the env.
     base_path: str = "s3://upstream-gridded-zarrs"
     k8s_secret_names: Sequence[str] = ["upstream-gridded-zarrs-key"]
+    format: DatasetFormat = DatasetFormat.ZARR3
 
 
 # Registry of all DynamicalDatasets.

src/reformatters/common/dynamical_dataset.py

@@ -11,7 +11,6 @@
 import sentry_sdk
 import typer
 import xarray as xr
-import zarr
 from pydantic import computed_field
 
 from reformatters.common import docker, template_utils, validation
@@ -27,13 +26,12 @@
 from reformatters.common.logging import get_logger
 from reformatters.common.pydantic import FrozenBaseModel
 from reformatters.common.region_job import RegionJob, SourceFileCoord
+from reformatters.common.storage import StorageConfig, StoreFactory
 from reformatters.common.template_config import TemplateConfig
 from reformatters.common.types import DatetimeLike
 from reformatters.common.zarr import (
     copy_zarr_metadata,
     get_local_tmp_store,
-    get_mode,
-    get_zarr_store,
 )
 
 DATA_VAR = TypeVar("DATA_VAR", bound=DataVar[Any])
@@ -42,20 +40,22 @@
 logger = get_logger(__name__)
 
 
-class DynamicalDatasetStorageConfig(FrozenBaseModel):
-    """Configuration for the storage of a dataset in production."""
-
-    base_path: str
-    k8s_secret_names: Sequence[str] = []
-
-
 class DynamicalDataset(FrozenBaseModel, Generic[DATA_VAR, SOURCE_FILE_COORD]):
     """Top level class managing a dataset configuration and processing."""
 
     template_config: TemplateConfig[DATA_VAR]
     region_job_class: type[RegionJob[DATA_VAR, SOURCE_FILE_COORD]]
 
-    storage_config: DynamicalDatasetStorageConfig
+    storage_config: StorageConfig
+
+    @computed_field  # type: ignore[prop-decorator]
+    @property
+    def primary_store_factory(self) -> StoreFactory:
+        return StoreFactory(
+            storage_config=self.storage_config,
+            dataset_id=self.dataset_id,
+            template_config_version=self.template_config.version,
+        )
 
     def operational_kubernetes_resources(self, image_tag: str) -> Iterable[CronJob]:
         """
@@ -127,27 +127,29 @@ def update(
     ) -> None:
         """Update an existing dataset with the latest data."""
         with self._monitor(ReformatCronJob, reformat_job_name):
-            final_store = self._final_store()
             tmp_store = self._tmp_store()
 
             jobs, template_ds = self.region_job_class.operational_update_jobs(
-                final_store=final_store,
+                primary_store_factory=self.primary_store_factory,
                 tmp_store=tmp_store,
                 get_template_fn=self._get_template,
                 append_dim=self.template_config.append_dim,
                 all_data_vars=self.template_config.data_vars,
                 reformat_job_name=reformat_job_name,
             )
-            template_utils.write_metadata(template_ds, tmp_store, get_mode(tmp_store))
+            template_utils.write_metadata(template_ds, tmp_store)
+
             for job in jobs:
                 process_results = job.process()
                 updated_template = job.update_template_with_results(process_results)
-                template_utils.write_metadata(
-                    updated_template, tmp_store, get_mode(tmp_store)
-                )
-                copy_zarr_metadata(updated_template, tmp_store, final_store)
+                # overwrite the tmp store metadata with updated template
+                template_utils.write_metadata(updated_template, tmp_store)
+                primary_store = self.primary_store_factory.store()
+                copy_zarr_metadata(updated_template, tmp_store, primary_store)
 
-        logger.info(f"Operational update complete. Wrote to store: {final_store}")
+        logger.info(
+            f"Operational update complete. Wrote to store: {self.primary_store_factory.store()}"
+        )
 
     def backfill_kubernetes(
         self,
@@ -164,15 +166,12 @@ def backfill_kubernetes(
         image_tag = docker_image or docker.build_and_push_image()
 
         template_ds = self._get_template(append_dim_end)
-        final_store = self._final_store()
-        logger.info(f"Writing zarr metadata to {final_store}")
-
-        template_utils.write_metadata(template_ds, final_store, get_mode(final_store))
+        template_utils.write_metadata(template_ds, self.primary_store_factory)
 
         num_jobs = len(
             self.region_job_class.get_jobs(
                 kind="backfill",
-                final_store=final_store,
+                primary_store_factory=self.primary_store_factory,
                 tmp_store=self._tmp_store(),
                 template_ds=template_ds,
                 append_dim=self.template_config.append_dim,
@@ -259,9 +258,7 @@ def backfill_local(
     ) -> None:
         """Run dataset reformatting locally in this process."""
         template_ds = self._get_template(append_dim_end)
-        final_store = self._final_store()
-
-        template_utils.write_metadata(template_ds, final_store, get_mode(final_store))
+        template_utils.write_metadata(template_ds, self.primary_store_factory)
 
         self.process_backfill_region_jobs(
             append_dim_end,
@@ -273,7 +270,7 @@ def backfill_local(
             filter_contains=filter_contains,
             filter_variable_names=filter_variable_names,
         )
-        logger.info(f"Done writing to {final_store}")
+        logger.info(f"Done writing to {self.primary_store_factory.store()}")
 
     def process_backfill_region_jobs(
         self,
@@ -291,7 +288,7 @@ def process_backfill_region_jobs(
 
         region_jobs = self.region_job_class.get_jobs(
             kind="backfill",
-            final_store=self._final_store(),
+            primary_store_factory=self.primary_store_factory,
             tmp_store=self._tmp_store(),
             template_ds=self._get_template(append_dim_end),
             append_dim=self.template_config.append_dim,
@@ -320,7 +317,7 @@ def validate_dataset(
     ) -> None:
         """Validate the dataset, raising an exception if it is invalid."""
         with self._monitor(ValidationCronJob, reformat_job_name):
-            store = self._final_store()
+            store = self.primary_store_factory.store()
             validation.validate_dataset(store, validators=self.validators())
 
         logger.info(f"Done validating {store}")
@@ -339,13 +336,6 @@ def get_cli(
         app.command("validate")(self.validate_dataset)
         return app
 
-    def _final_store(self) -> zarr.abc.store.Store:
-        return get_zarr_store(
-            self.storage_config.base_path,
-            self.template_config.dataset_id,
-            self.template_config.version,
-        )
-
     def _tmp_store(self) -> Path:
         return get_local_tmp_store()
 

src/reformatters/common/region_job.py

@@ -15,7 +15,6 @@
 import pandas as pd
 import pydantic
 import xarray as xr
-import zarr
 from pydantic import AfterValidator, Field, computed_field
 
 from reformatters.common import template_utils
@@ -28,6 +27,7 @@
     create_data_array_and_template,
 )
 from reformatters.common.shared_memory_utils import make_shared_buffer, write_shards
+from reformatters.common.storage import StoreFactory
 from reformatters.common.types import (
     AppendDim,
     ArrayND,
@@ -36,7 +36,7 @@
     Timestamp,
 )
 from reformatters.common.update_progress_tracker import UpdateProgressTracker
-from reformatters.common.zarr import copy_data_var, get_mode
+from reformatters.common.zarr import copy_data_var
 
 log = get_logger(__name__)
 
@@ -101,7 +101,7 @@ def region_slice(s: slice) -> slice:
 
 
 class RegionJob(pydantic.BaseModel, Generic[DATA_VAR, SOURCE_FILE_COORD]):
-    final_store: zarr.abc.store.Store
+    primary_store_factory: StoreFactory
     tmp_store: Path
     template_ds: xr.Dataset
     data_vars: Sequence[DATA_VAR]
@@ -228,7 +228,7 @@ def update_template_with_results(
         Subclasses should implement this method to apply dataset-specific adjustments
         based on the processing results. Examples include:
         - Trimming dataset along append_dim to only include successfully processed data
-        - Loading existing coordinate values from final_store and updating them based on results
+        - Loading existing coordinate values from the primary store and updating them based on results
         - Updating metadata based on what was actually processed vs what was planned
 
         The default implementation here trims along append_dim to end at the most recent
@@ -266,7 +266,7 @@ def update_template_with_results(
     @classmethod
     def operational_update_jobs(
         cls,
-        final_store: zarr.abc.store.Store,
+        primary_store_factory: StoreFactory,
         tmp_store: Path,
         get_template_fn: Callable[[DatetimeLike], xr.Dataset],
         append_dim: AppendDim,
@@ -284,16 +284,16 @@ def operational_update_jobs(
 
         The exact logic is dataset-specific, but it generally follows this pattern:
         1. Figure out the range of time to process: append_dim_start (inclusive) and append_dim_end (exclusive)
-            a. Read existing data from final_store to determine what's already processed
+            a. Read existing data from the primary store to determine what's already processed
             b. Optionally identify recent incomplete/non-final data for reprocessing
         2. Call get_template_fn(append_dim_end) to get the template_ds
         3. Create RegionJob instances by calling cls.get_jobs(..., filter_start=append_dim_start)
 
         Parameters
         ----------
-        final_store : zarr.abc.store.Store
-            The destination Zarr store to read existing data from and write updates to.
-        tmp_store : zarr.abc.store.Store | Path
+        primary_store_factory : StoreFactory
+            The factory to get the primary store to read existing data from and write updates to.
+        tmp_store : Path
             The temporary Zarr store to write into while processing.
         get_template_fn : Callable[[DatetimeLike], xr.Dataset]
             Function to get the template_ds for the operational update.
@@ -331,7 +331,7 @@ def dataset_id(self) -> str:
     def get_jobs(
         cls,
         kind: Literal["backfill", "operational-update"],
-        final_store: zarr.abc.store.Store,
+        primary_store_factory: StoreFactory,
         tmp_store: Path,
         template_ds: xr.Dataset,
         append_dim: AppendDim,
@@ -357,9 +357,9 @@ def get_jobs(
 
         Parameters
         ----------
-        final_store : zarr.abc.store.Store
-            The destination Zarr store to write into.
-        tmp_store : zarr.abc.store.Store | Path
+        primary_store_factory : StoreFactory
+            The factory to get the primary store to read existing data from and write updates to.
+        tmp_store : Path
             The temporary Zarr store to write into while processing.
         template_ds : xr.Dataset
             Dataset template defining structure and metadata.
@@ -442,7 +442,7 @@ def get_jobs(
 
         all_jobs = [
             cls(
-                final_store=final_store,
+                primary_store_factory=primary_store_factory,
                 tmp_store=tmp_store,
                 template_ds=template_ds,
                 data_vars=data_var_group,
@@ -468,7 +468,7 @@ def process(self) -> Mapping[str, Sequence[SOURCE_FILE_COORD]]:
                 i.   Read data from source files into the shared array
                 ii.  Apply any required data transformations (e.g., rounding, deaccumulation)
                 iii. Write output shards to the tmp_store
-                iv.  Upload chunk data from tmp_store to final_store
+                iv.  Upload chunk data from tmp_store to the primary store
 
         Returns
         -------
@@ -477,8 +477,10 @@ def process(self) -> Mapping[str, Sequence[SOURCE_FILE_COORD]]:
         """
         processing_region_ds, output_region_ds = self._get_region_datasets()
 
+        primary_store = self.primary_store_factory.store()
+
         progress_tracker = UpdateProgressTracker(
-            self.final_store, self.reformat_job_name, self.region.start
+            primary_store, self.reformat_job_name, self.region.start
         )
         data_vars_to_process: Sequence[DATA_VAR] = progress_tracker.get_unprocessed(
             self.data_vars
@@ -489,9 +491,7 @@ def process(self) -> Mapping[str, Sequence[SOURCE_FILE_COORD]]:
                 data_var_groups, self.max_vars_per_download_group
             )
 
-        template_utils.write_metadata(
-            self.template_ds, self.tmp_store, get_mode(self.tmp_store)
-        )
+        template_utils.write_metadata(self.template_ds, self.tmp_store)
 
         results: dict[str, Sequence[SOURCE_FILE_COORD]] = {}
         upload_futures: list[Any] = []
@@ -555,7 +555,7 @@ def process(self) -> Mapping[str, Sequence[SOURCE_FILE_COORD]]:
                             self.template_ds,
                             self.append_dim,
                             self.tmp_store,
-                            self.final_store,
+                            primary_store,
                             partial(progress_tracker.record_completion, data_var.name),
                         )
                     )