Skip to content

Expose resolved blob versions #116

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
damonbayer wants to merge 1 commit into from
Closed

Expose resolved blob versions #116

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
143 changes: 98 additions & 45 deletions cfa/dataops/catalog.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -440,13 +488,15 @@ 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(
self,
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

Expand All @@ -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
Expand All @@ -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(
Expand Down
8 changes: 8 additions & 0 deletions cfa/dataops/stub_templates/catalog.pyi.mako
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -59,20 +64,23 @@ 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(
self,
output: Literal["polars", "pl"],
version_spec: str | None = None,
selection: Literal["newest", "oldest"] = "newest",
print_version: bool = True,
) -> pl.DataFrame: ...
@overload
def get_dataframe(
self,
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(
Expand Down
8 changes: 8 additions & 0 deletions tests/fixtures/expected_catalog_stub.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -59,20 +64,23 @@ 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(
self,
output: Literal["polars", "pl"],
version_spec: str | None = None,
selection: Literal["newest", "oldest"] = "newest",
print_version: bool = True,
) -> pl.DataFrame: ...
@overload
def get_dataframe(
self,
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(
Expand Down
Loading
Loading