warpzone-sdk 14.1.0__tar.gz → 14.3.0__tar.gz

{warpzone_sdk-14.1.0 → warpzone_sdk-14.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: warpzone-sdk
-Version: 14.1.0
+Version: 14.3.0
 Summary: The main objective of this package is to centralize logic used to interact with Azure Functions, Azure Service Bus and Azure Table Storage
 Author: Team Enigma
 Author-email: enigma@energinet.dk
@@ -22,7 +22,7 @@ Requires-Dist: azure-servicebus (>=7.8.0,<7.9.0)
 Requires-Dist: azure-storage-blob (>=12.14.1)
 Requires-Dist: cryptography (==43.0.3)
 Requires-Dist: datamazing (>=5.1.6)
-Requires-Dist: deltalake (>=1.3.0)
+Requires-Dist: deltalake (==1.2.1)
 Requires-Dist: numpy (>=1.26.4)
 Requires-Dist: obstore (>=0.8.2)
 Requires-Dist: opentelemetry-sdk (>=1.32.0)

{warpzone_sdk-14.1.0 → warpzone_sdk-14.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "warpzone-sdk"
-version = "14.1.0"
+version = "14.3.0"
 description = "The main objective of this package is to centralize logic used to interact with Azure Functions, Azure Service Bus and Azure Table Storage"
 authors = [{ name = "Team Enigma", email = "enigma@energinet.dk" }]
 requires-python = ">=3.10"
@@ -14,7 +14,7 @@ requires-poetry = ">=2.2"
 pandas = ">=2.0.3"
 polars = ">=1.33.1"
 obstore = ">=0.8.2"
-deltalake = ">=1.3.0"
+deltalake = "==1.2.1" # pin to avoid breaking changes in 1.3.0. Follow Github issue here https://github.com/delta-io/delta-rs/issues/3939
 pyarrow = ">=19.0.0"
 datamazing = ">=5.1.6"
 azure-core = ">=1.26.3"

{warpzone_sdk-14.1.0 → warpzone_sdk-14.3.0}/warpzone/db/client.py

@@ -32,29 +32,51 @@ class WarpzoneDatabaseClient:
         self,
         path: str,
         storage_options: dict[str, str] | None = None,
+        table_prefix: str = "",
     ):
         self.store = Store(
             path=path,
             storage_options=storage_options,
         )
+        self.table_prefix = table_prefix
 
     @classmethod
-    def from_storage_account(
+    def from_resource_name(
         cls,
         storage_account: str,
         container_name: str = "datasets",
+        sub_path: str = "",
+        table_prefix: str = "",
         credential: (
             AzureNamedKeyCredential | AzureSasCredential | TokenCredential
         ) = DefaultAzureCredential(),
     ):
+        """Create a WarpzoneDatabaseClient from resource name (storage account).
+        This assumes the path of the delta lake is of the form:
+        abfss://{container_name}@{storage_account}.dfs.core.windows.net/{sub_path}
+
+        Args:
+            storage_account (str): Storage account name.
+            container_name (str, optional): Container name. Defaults to "datasets".
+            sub_path (str, optional): Sub-path within the container. Defaults to "".
+            table_prefix (str, optional): Table prefix to use (e.g. `mz_` for archive).
+            Defaults to "".
+            credential (optional): Azure credential to use.
+            Defaults to DefaultAzureCredential().
+        """
         path = f"abfss://{container_name}@{storage_account}.dfs.core.windows.net"
+        if sub_path:
+            path += f"/{sub_path}"
+
         token = credential.get_token("https://storage.azure.com/.default")
         storage_options = {
             "account_name": storage_account,
             "token": token.token,
         }
 
-        return cls(path=path, storage_options=storage_options)
+        return cls(
+            path=path, storage_options=storage_options, table_prefix=table_prefix
+        )
 
     def get_unit_and_multiple(self, timedelta: pd.Timedelta) -> tuple[str | None, int]:
         """
@@ -148,7 +170,35 @@ class WarpzoneDatabaseClient:
         time_interval: Optional[pdz.TimeInterval] = None,
         time_travel: Optional[pdz.TimeTravel] = None,
         filters: Optional[dict[str, object]] = None,
+        include_validity_period_columns: bool = False,
+        include_generated_columns: bool = False,
     ) -> pd.DataFrame:
+        """Query table.
+        Query defaults are set to match old Table Storage client behavior.
+        Time travel defaults to "as of now"
+        Validity period columns are dropped by default.
+        Generated columns are dropped by default.
+
+        Args:
+            table_name (str): Name of the table
+            time_interval (Optional[pdz.TimeInterval], optional): Time interval for the
+            query. Defaults to None.
+            time_travel (Optional[pdz.TimeTravel], optional): Time travel information.
+            Defaults to None.
+            filters (Optional[dict[str, object]], optional): Filters to apply to the
+            query.
+            Defaults to None.
+            include_validity_period_columns (bool, optional): Whether to include
+            validity period columns in the result;
+            (e.g. `valid_from_time_utc`, `valid_to_time_utc`).
+            Defaults to False.
+            include_generated_columns (bool, optional): Whether to include generated
+            columns in the result; (e.g. `valid_from_time_utc`, `valid_to_time_utc`).
+            Defaults to False.
+
+        Returns:
+            pd.DataFrame: The result of the query.
+        """
         # We do 'camelCaseToSnake_case' conversion here because the old
         # naming convention used in WarpZone was CamelCase, while the new
         # naming convention is snake_case. The goal is to remove this
@@ -191,16 +241,16 @@ class WarpzoneDatabaseClient:
         for col in pd_df.select_dtypes(include=["datetime", "datetimetz"]).columns:
             pd_df[col] = pd_df[col].dt.floor("s").dt.as_unit("ns")
 
-        # We remove the valid-from and valid-to columns, as well
-        # as any generated columns, as this was not present
-        # in the old solution (Azure Table Stroage)
-        generated_cols = []
-        for field in table.schema().fields:
-            if field.generated_as is not None:
-                generated_cols.append(field.column_name)
-
-        pd_df = pd_df.drop(
-            columns=["valid_from_time_utc", "valid_to_time_utc"] + generated_cols
-        )
+        # Drop generated columns
+        if not include_generated_columns:
+            generated_cols = []
+            for field in table.schema().fields:
+                if field.generated_as is not None:
+                    generated_cols.append(field.column_name)
+            pd_df = pd_df.drop(columns=generated_cols)
+
+        # Drop valid-from/to columns
+        if not include_validity_period_columns:
+            pd_df = pd_df.drop(columns=["valid_from_time_utc", "valid_to_time_utc"])
 
         return pd_df

{warpzone_sdk-14.1.0 → warpzone_sdk-14.3.0}/warpzone/function/integrations.py

@@ -62,7 +62,7 @@ def get_db_client() -> WarpzoneDatabaseClient:
 
 
 def get_delta_db_client() -> WarpzoneDeltaDatabaseClient:
-    db_client = WarpzoneDeltaDatabaseClient.from_storage_account(
+    db_client = WarpzoneDeltaDatabaseClient.from_resource_name(
         os.environ["OPERATIONAL_DATA_STORAGE_ACCOUNT"],
         credential=_credential,
     )