diff --git a/cfa/dataops/catalog.py b/cfa/dataops/catalog.py index 6398b51..32eb414 100644 --- a/cfa/dataops/catalog.py +++ b/cfa/dataops/catalog.py @@ -5,6 +5,7 @@ import pkgutil from collections.abc import Sequence from configparser import ConfigParser +from dataclasses import dataclass from importlib import import_module from io import BytesIO from pathlib import PurePosixPath @@ -88,6 +89,12 @@ class CatalogNamespace(SimpleNamespace): """Runtime namespace wrapper for catalog access.""" +@dataclass(frozen=True) +class _ResolvedVersion: + version: str | list[str] | None + blob_prefixes: list[str] + + class DatasetEndpoint: """The DatasetEndpoint class for including in the datacat namespace. This ends the namespace branching at a config file and creates all the @@ -304,6 +311,71 @@ def get_file_ext( version_spec=version_spec, selection=selection, print_version=False )[0]["name"].split(".")[-1] + def resolve_version( + self, + version_spec: str | None = None, + selection: Literal["newest", "oldest", "all"] = "newest", + ) -> dict[str, Any]: + """Resolve a version specifier without loading blob data. + + Args: + version_spec (str | None, optional): Version specifier to pass through to + ``version_matcher``. Defaults to ``None``. + selection (Literal["newest", "oldest", "all"], optional): When matching multiple versions, choose the + newest matching version when ``"newest"``, the oldest matching version + when ``"oldest"``, or all matching versions when ``"all"``. + + Returns: + dict[str, Any]: Metadata for the resolved version, including the + selected version, original version specifier, selection strategy, + resolved blob prefixes, and Azure blob prefix URLs. + + Raises: + ValueError: If the requested version cannot be resolved. + """ + resolved = self._resolve_version( + version_spec=version_spec, selection=selection + ) + + return { + "version": resolved.version, + "version_spec": version_spec, + "selection": selection, + "account": self.account, + "container": self.container, + "blob_prefixes": resolved.blob_prefixes, + "blob_urls": [ + f"az://{self.container}/{blob_prefix}" + for blob_prefix in resolved.blob_prefixes + ], + } + + def _resolve_version( + self, + version_spec: str | None = None, + selection: Literal["newest", "oldest", "all"] = "newest", + ) -> _ResolvedVersion: + """Resolve a version specifier into version metadata.""" + if not check_ext_env(): + raise RuntimeError("No EXT access configured.") + if self.is_ledger: + return _ResolvedVersion( + version=None, + blob_prefixes=[f"{self.prefix.removesuffix('/')}/"], + ) + + available_versions = self.get_versions() + version = version_matcher(version_spec, available_versions, selection=selection) + if not version: + raise ValueError( + f"Version {version} not found in available versions: {available_versions}" + ) + if isinstance(version, list): + blob_prefixes = [f"{self.prefix}/{v}/" for v in version] + else: + blob_prefixes = [f"{self.prefix}/{version}/"] + return _ResolvedVersion(version=version, blob_prefixes=blob_prefixes) + def _get_version_blobs( self, version_spec: str | None = None, @@ -328,51 +400,25 @@ def _get_version_blobs( Raises: ValueError: If the requested version cannot be resolved. """ - # check credential access - if not check_ext_env(): - raise RuntimeError("No EXT access configured.") - if not self.is_ledger: - available_versions = self.get_versions() - version = version_matcher( - version_spec, available_versions, selection=selection - ) - if not version: - raise ValueError( - f"Version {version} not found in available versions: {available_versions}" - ) - if print_version: - print(f"Using version: {version}") - if isinstance(version, list): - walk_path = [f"{self.prefix}/{v}/" for v in version] - else: - walk_path = f"{self.prefix}/{version}/" - else: - walk_path = f"{self.prefix.removesuffix('/')}/" - if isinstance(walk_path, list): - all_blobs = [] - for wp in walk_path: - all_blobs.extend( - walk_blobs_in_container( - name_starts_with=wp, - account_name=self.account, - container_name=self.container, - ) + resolved = self._resolve_version( + version_spec=version_spec, selection=selection + ) + if print_version and not self.is_ledger: + print(f"Using version: {resolved.version}") + + all_blobs = [] + for blob_prefix in resolved.blob_prefixes: + all_blobs.extend( + walk_blobs_in_container( + name_starts_with=blob_prefix, + account_name=self.account, + container_name=self.container, ) - return sorted( - all_blobs, - key=lambda x: x["creation_time"], - ) - else: - return sorted( - list( - walk_blobs_in_container( - name_starts_with=walk_path, - account_name=self.account, - container_name=self.container, - ) - ), - key=lambda x: x["creation_time"], ) + return sorted( + all_blobs, + key=lambda x: x["creation_time"], + ) def download_version_to_local( self, @@ -424,6 +470,7 @@ def get_dataframe( output: Literal["pandas", "pd"] = "pandas", version_spec: str | None = None, selection: Literal["newest", "oldest"] = "newest", + print_version: bool = True, ) -> pd.DataFrame: ... @overload @@ -432,6 +479,7 @@ def get_dataframe( output: Literal["polars", "pl"], version_spec: str | None = None, selection: Literal["newest", "oldest"] = "newest", + print_version: bool = True, ) -> pl.DataFrame: ... @overload @@ -440,6 +488,7 @@ def get_dataframe( output: Literal["pl_lazy", "lazy"], version_spec: str | None = None, selection: Literal["newest", "oldest"] = "newest", + print_version: bool = True, ) -> pl.LazyFrame: ... def get_dataframe( @@ -447,6 +496,7 @@ def get_dataframe( output: Literal["pandas", "pd", "polars", "pl", "pl_lazy", "lazy"] = "pandas", version_spec: str | None = None, selection: Literal["newest", "oldest"] = "newest", + print_version: bool = True, ) -> pd.DataFrame | pl.DataFrame | pl.LazyFrame: """Get the data as a pandas or polars dataframe @@ -455,7 +505,8 @@ def get_dataframe( either 'pandas' or 'polars' or 'pl_lazy'. Defaults to "pandas". version_spec (str, optional): the version of the data to get. Defaults to "latest". - selection (Literal["newest", "oldest", "all"], optional): whether to get the newest, oldest, or all matching versions. Defaults to "newest". + selection (Literal["newest", "oldest"], optional): whether to get the newest or oldest matching version. Defaults to "newest". + print_version (bool, optional): whether to print the resolved version. Defaults to True. Raises: ValueError: if output is not one of @@ -472,7 +523,9 @@ def get_dataframe( ) # Fetch version blobs once and validate before deriving file extension. version_blobs = self._get_version_blobs( - version_spec=version_spec, selection=selection + version_spec=version_spec, + selection=selection, + print_version=print_version, ) if not version_blobs: raise ValueError( diff --git a/cfa/dataops/stub_templates/catalog.pyi.mako b/cfa/dataops/stub_templates/catalog.pyi.mako index fffda37..a267b80 100644 --- a/cfa/dataops/stub_templates/catalog.pyi.mako +++ b/cfa/dataops/stub_templates/catalog.pyi.mako @@ -46,6 +46,11 @@ class BlobEndpoint: version_spec: str | None = None, selection: Literal["newest", "oldest", "all"] = "newest", ) -> str: ... + def resolve_version( + self, + version_spec: str | None = None, + selection: Literal["newest", "oldest", "all"] = "newest", + ) -> dict[str, Any]: ... def download_version_to_local( self, local_path: str, @@ -59,6 +64,7 @@ class BlobEndpoint: output: Literal["pandas", "pd"] = "pandas", version_spec: str | None = None, selection: Literal["newest", "oldest"] = "newest", + print_version: bool = True, ) -> pd.DataFrame: ... @overload def get_dataframe( @@ -66,6 +72,7 @@ class BlobEndpoint: output: Literal["polars", "pl"], version_spec: str | None = None, selection: Literal["newest", "oldest"] = "newest", + print_version: bool = True, ) -> pl.DataFrame: ... @overload def get_dataframe( @@ -73,6 +80,7 @@ class BlobEndpoint: output: Literal["pl_lazy", "lazy"], version_spec: str | None = None, selection: Literal["newest", "oldest"] = "newest", + print_version: bool = True, ) -> pl.LazyFrame: ... def ledger_entry(self, action: str) -> None: ... def save_dataframe( diff --git a/tests/fixtures/expected_catalog_stub.pyi b/tests/fixtures/expected_catalog_stub.pyi index 04a6cdf..4711673 100644 --- a/tests/fixtures/expected_catalog_stub.pyi +++ b/tests/fixtures/expected_catalog_stub.pyi @@ -46,6 +46,11 @@ class BlobEndpoint: version_spec: str | None = None, selection: Literal["newest", "oldest", "all"] = "newest", ) -> str: ... + def resolve_version( + self, + version_spec: str | None = None, + selection: Literal["newest", "oldest", "all"] = "newest", + ) -> dict[str, Any]: ... def download_version_to_local( self, local_path: str, @@ -59,6 +64,7 @@ class BlobEndpoint: output: Literal["pandas", "pd"] = "pandas", version_spec: str | None = None, selection: Literal["newest", "oldest"] = "newest", + print_version: bool = True, ) -> pd.DataFrame: ... @overload def get_dataframe( @@ -66,6 +72,7 @@ class BlobEndpoint: output: Literal["polars", "pl"], version_spec: str | None = None, selection: Literal["newest", "oldest"] = "newest", + print_version: bool = True, ) -> pl.DataFrame: ... @overload def get_dataframe( @@ -73,6 +80,7 @@ class BlobEndpoint: output: Literal["pl_lazy", "lazy"], version_spec: str | None = None, selection: Literal["newest", "oldest"] = "newest", + print_version: bool = True, ) -> pl.LazyFrame: ... def ledger_entry(self, action: str) -> None: ... def save_dataframe( diff --git a/tests/test_blob_endpoint_resolve_version.py b/tests/test_blob_endpoint_resolve_version.py new file mode 100644 index 0000000..a602aa2 --- /dev/null +++ b/tests/test_blob_endpoint_resolve_version.py @@ -0,0 +1,148 @@ +"""Tests for BlobEndpoint.resolve_version.""" + +import pytest + +from cfa.dataops.catalog import BlobEndpoint + + +@pytest.fixture +def blob_endpoint(): + ledger_location = { + "account": "account_test", + "container": "container_test", + "prefix": "_access/test/ledger/", + } + return BlobEndpoint( + account="account_test", + container="container_test", + prefix="test/prefix", + ledger_location=ledger_location, + ns="test.endpoint", + ) + + +def test_resolve_version_returns_metadata_without_loading_blob_data( + mocker, blob_endpoint +): + mocker.patch("cfa.dataops.catalog.check_ext_env", return_value=True) + mock_get_versions = mocker.patch.object( + blob_endpoint, + "get_versions", + return_value=["2026-07-08T00-00-00", "2026-07-07T00-00-00"], + ) + mock_get_version_blobs = mocker.patch.object(blob_endpoint, "_get_version_blobs") + mock_walk_blobs = mocker.patch("cfa.dataops.catalog.walk_blobs_in_container") + mock_read_blob_stream = mocker.patch("cfa.dataops.catalog.read_blob_stream") + + metadata = blob_endpoint.resolve_version(version_spec="<=2026-07-09T00-00-00") + + assert metadata == { + "version": "2026-07-08T00-00-00", + "version_spec": "<=2026-07-09T00-00-00", + "selection": "newest", + "account": "account_test", + "container": "container_test", + "blob_prefixes": ["test/prefix/2026-07-08T00-00-00/"], + "blob_urls": ["az://container_test/test/prefix/2026-07-08T00-00-00/"], + } + mock_get_versions.assert_called_once_with() + mock_get_version_blobs.assert_not_called() + mock_walk_blobs.assert_not_called() + mock_read_blob_stream.assert_not_called() + + +def test_resolve_version_returns_all_matching_versions(mocker, blob_endpoint): + mocker.patch("cfa.dataops.catalog.check_ext_env", return_value=True) + mocker.patch.object( + blob_endpoint, + "get_versions", + return_value=[ + "2026-07-10T00-00-00", + "2026-07-08T00-00-00", + "2026-07-07T00-00-00", + ], + ) + + metadata = blob_endpoint.resolve_version( + version_spec="<=2026-07-09T00-00-00", selection="all" + ) + + assert metadata["version"] == [ + "2026-07-08T00-00-00", + "2026-07-07T00-00-00", + ] + assert metadata["blob_prefixes"] == [ + "test/prefix/2026-07-08T00-00-00/", + "test/prefix/2026-07-07T00-00-00/", + ] + assert metadata["blob_urls"] == [ + "az://container_test/test/prefix/2026-07-08T00-00-00/", + "az://container_test/test/prefix/2026-07-07T00-00-00/", + ] + + +def test_get_version_blobs_reuses_resolved_version_paths(mocker, blob_endpoint): + mocker.patch("cfa.dataops.catalog.check_ext_env", return_value=True) + mocker.patch.object( + blob_endpoint, + "get_versions", + return_value=["2026-07-08T00-00-00"], + ) + mock_walk_blobs = mocker.patch( + "cfa.dataops.catalog.walk_blobs_in_container", + return_value=[ + { + "name": "test/prefix/2026-07-08T00-00-00/data.parquet", + "creation_time": "2026-07-08T00:00:00", + } + ], + ) + + blobs = blob_endpoint._get_version_blobs( + version_spec="<=2026-07-09T00-00-00", print_version=False + ) + + assert blobs == [ + { + "name": "test/prefix/2026-07-08T00-00-00/data.parquet", + "creation_time": "2026-07-08T00:00:00", + } + ] + mock_walk_blobs.assert_called_once_with( + name_starts_with="test/prefix/2026-07-08T00-00-00/", + account_name="account_test", + container_name="container_test", + ) + + +def test_get_dataframe_can_suppress_version_print(mocker, blob_endpoint): + mocker.patch("cfa.dataops.catalog.check_ext_env", return_value=True) + mock_get_version_blobs = mocker.patch.object( + blob_endpoint, + "_get_version_blobs", + return_value=[ + { + "name": "test/prefix/2026-07-08T00-00-00/data.csv", + "creation_time": "2026-07-08T00:00:00", + } + ], + ) + mock_read_blobs = mocker.patch.object( + blob_endpoint, + "read_blobs", + return_value=[b"value\n1\n"], + ) + + df = blob_endpoint.get_dataframe(print_version=False) + + assert df["value"].tolist() == [1] + mock_get_version_blobs.assert_called_once_with( + version_spec=None, + selection="newest", + print_version=False, + ) + mock_read_blobs.assert_called_once_with( + version_spec=None, + selection="newest", + print_version=False, + ) diff --git a/typings/cfa/dataops/catalog.pyi b/typings/cfa/dataops/catalog.pyi index 4681916..587f960 100644 --- a/typings/cfa/dataops/catalog.pyi +++ b/typings/cfa/dataops/catalog.pyi @@ -44,6 +44,11 @@ class BlobEndpoint: version_spec: str | None = None, selection: Literal["newest", "oldest", "all"] = "newest", ) -> str: ... + def resolve_version( + self, + version_spec: str | None = None, + selection: Literal["newest", "oldest", "all"] = "newest", + ) -> dict[str, Any]: ... def download_version_to_local( self, local_path: str, @@ -57,6 +62,7 @@ class BlobEndpoint: output: Literal["pandas", "pd"] = "pandas", version_spec: str | None = None, selection: Literal["newest", "oldest"] = "newest", + print_version: bool = True, ) -> pd.DataFrame: ... @overload def get_dataframe( @@ -64,6 +70,7 @@ class BlobEndpoint: output: Literal["polars", "pl"], version_spec: str | None = None, selection: Literal["newest", "oldest"] = "newest", + print_version: bool = True, ) -> pl.DataFrame: ... @overload def get_dataframe( @@ -71,6 +78,7 @@ class BlobEndpoint: output: Literal["pl_lazy", "lazy"], version_spec: str | None = None, selection: Literal["newest", "oldest"] = "newest", + print_version: bool = True, ) -> pl.LazyFrame: ... def ledger_entry(self, action: str) -> None: ... def save_dataframe(