warpzone-sdk 15.0.0.dev14__py3-none-any.whl → 15.1.0__py3-none-any.whl

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,41 @@ class WarpzoneDatabaseClient:
         time_interval: Optional[pdz.TimeInterval] = None,
         time_travel: Optional[pdz.TimeTravel] = None,
         filters: Optional[dict[str, object]] = None,
+        columns: Optional[list[str]] = 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.
+            columns (Optional[list[str]], optional): Columns to return.
+            Selecting columns can significantly improve query performance.
+            Defaults to None, meaning all columns will be returned.
+            include_validity_period_columns (bool, optional): Whether to include
+            validity period columns in the result;
+            (`valid_from_time_utc`, `valid_to_time_utc`).
+            Defaults to False. If set to True while using `columns`-argument, make sure
+            to include these columns in the `columns`-list.
+            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. If set to True while using `columns`-argument, make sure
+            to include these columns in the `columns`-list.
+
+        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
@@ -164,7 +220,6 @@ class WarpzoneDatabaseClient:
                     hyper_slice.append((key, "in", value))
                 else:
                     hyper_slice.append((key, "=", value))
-
         if time_interval:
             hyper_slice.append(("time_utc", ">=", time_interval.left))
             hyper_slice.append(("time_utc", "<=", time_interval.right))
@@ -180,9 +235,9 @@ class WarpzoneDatabaseClient:
             valid_from_column="valid_from_time_utc",
             valid_to_column="valid_to_time_utc",
         )
-        hyper_slice.extend(tt_filter)
 
-        pl_df = table.read(hyper_slice=HyperSlice(hyper_slice))
+        hyper_slice.extend(tt_filter)
+        pl_df = table.read(hyper_slice=HyperSlice(hyper_slice), columns=columns)
 
         pd_df = pl_df.to_pandas()
 
@@ -191,16 +246,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/deltastorage/store.py

@@ -2,6 +2,7 @@ import os
 
 import deltalake as dl
 import obstore as obs
+from deltalake.exceptions import TableNotFoundError as DeltaTableNotFoundError
 
 from .schema import Schema
 from .table import Table
@@ -76,10 +77,19 @@ class Store:
         Args:
             table_name (str): Table name
         """
-        return dl.DeltaTable.is_deltatable(
-            table_uri=self._get_table_uri(table_name),
-            storage_options=self.storage_options,
-        )
+        # For some reason `deltalake.DeltaTable.is_deltatable()` can be very slow.
+        # deltalake has an issue open about this:
+        # https://github.com/delta-io/delta-rs/issues/3942
+        # For now we catch the exception when trying to load the table
+        try:
+            _ = dl.DeltaTable(
+                table_uri=self._get_table_uri(table_name),
+                storage_options=self.storage_options,
+                without_files=True,
+            )
+        except DeltaTableNotFoundError:
+            return False
+        return True
 
     def create_table(
         self,

warpzone/deltastorage/table.py

@@ -42,6 +42,7 @@ class Table:
         self.storage_options = storage_options
 
         self.table_name = self.table_uri.split("/")[-1]
+        self._delta_table = None
 
     def __repr__(self):
         return f"Table('{self.table_name}')"
@@ -49,12 +50,21 @@ class Table:
     @property
     def delta_table(self) -> dl.DeltaTable:
         """Get the Delta table object.
-        This is lazily computed to avoid using the same instance.
-        This is *important* because using the same instance can lead to transaction
+        As the `Table`-class is lazily initialized,
+        the `delta_table`-property is initialized on the first access
+        and saved for future use to minimize overhead.
+        It is *important* that this property is only initialized within
+        a lock when doing concurrent reads/writes
+        and not initialized when creating the `Table`-object.
+        This is important because using the same instance can lead to transaction
         issues in delta as DeltaTable uses metadata (transaction id) from
         the first time the object is instantiated.
         """
-        return dl.DeltaTable(self.table_uri, storage_options=self.storage_options)
+        if self._delta_table is None:
+            self._delta_table = dl.DeltaTable(
+                self.table_uri, storage_options=self.storage_options
+            )
+        return self._delta_table
 
     def partition_cols(self) -> list[str]:
         """Get the partition columns of the table"""
@@ -65,7 +75,9 @@ class Table:
         pa_schema = pa.schema(self.delta_table.schema())
         return Schema.from_arrow(pa_schema)
 
-    def read(self, hyper_slice: Optional[HyperSlice] = None) -> pl.DataFrame:
+    def read(
+        self, hyper_slice: Optional[HyperSlice] = None, columns=None
+    ) -> pl.DataFrame:
         """Read from Delta table
 
         Args:
@@ -88,6 +100,7 @@ class Table:
             partition_filters = [f for f in hyper_slice if f[0] in partition_cols]
 
         pyarrow_table_existing_data = delta_table.to_pyarrow_table(
+            columns=columns,
             partitions=partition_filters,
             filters=file_filters,
         )

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,
     )

warpzone/monitor/logs.py

@@ -34,6 +34,7 @@ LOGGING_IS_CONFIGURED = False
 
 def configure_logging():
     global LOGGING_IS_CONFIGURED
+    # Add thread locking to avoid race conditions during setup
     with _LOGGING_LOCK:
         if LOGGING_IS_CONFIGURED:
             # logging should only be set up once

warpzone/monitor/traces.py

@@ -23,6 +23,7 @@ TRACING_IS_CONFIGURED = False
 
 def configure_tracing():
     global TRACING_IS_CONFIGURED
+    # Add thread locking to avoid race conditions during setup
     with _TRACING_LOCK:
         if TRACING_IS_CONFIGURED:
             # tracing should only be set up once

warpzone/tablestorage/db/client.py

@@ -12,13 +12,10 @@ from azure.identity import DefaultAzureCredential
 
 from warpzone.blobstorage.client import WarpzoneBlobClient
 from warpzone.healthchecks import HealthCheckResult, check_health_of
-from warpzone.monitor import traces
 from warpzone.tablestorage.db import base_client
 from warpzone.tablestorage.db.table_config import DataType, TableMetadata
 from warpzone.tablestorage.tables.client import WarpzoneTableClient
 
-tracer = traces.get_tracer(__name__)
-
 
 class WarpzoneDatabaseClient:
     """Class to interact with Azure Table Storage for database queries

{warpzone_sdk-15.0.0.dev14.dist-info → warpzone_sdk-15.1.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: warpzone-sdk
-Version: 15.0.0.dev14
+Version: 15.1.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

{warpzone_sdk-15.0.0.dev14.dist-info → warpzone_sdk-15.1.0.dist-info}/RECORD

@@ -2,21 +2,21 @@ warpzone/__init__.py,sha256=Ay7znIadokCdHHqsPlWAqhdg-zdWhXHNtxVinViYm7o,1533
 warpzone/blobstorage/__init__.py,sha256=lnc0uiaGLF0qMi_rWhCpRSFvaj0CJEiMCAl6Yqn1ZiA,21
 warpzone/blobstorage/client.py,sha256=YwDV83acoCeHS_D_ydsTYwnf56rSTy9CKpsxqeoXmBs,4638
 warpzone/db/__init__.py,sha256=lnc0uiaGLF0qMi_rWhCpRSFvaj0CJEiMCAl6Yqn1ZiA,21
-warpzone/db/client.py,sha256=_LYIZIYwaii9dyTwpqoOCYbqzzcbJNgflzSvUqeTto8,6919
+warpzone/db/client.py,sha256=tqY_ANThnO72rBaB4sydigLhlSFwCmadg_6rVZwjO8M,9688
 warpzone/deltastorage/__init__.py,sha256=cV8sGT2N_N5Z-E179NMW5O7q3FUDrb3j5f-yVNlNPv0,152
 warpzone/deltastorage/data_types.py,sha256=tWjLO_0ig7-tYxSHShvrd0znA7FFDnS-wBuFClUQG2U,2059
 warpzone/deltastorage/generated_columns.py,sha256=Dr_bihM7v9JKCgBXxc3JQC2P2mUGobXsReRKOl6jDO4,4765
 warpzone/deltastorage/lock_client.py,sha256=z19ub2ahsUHdcriRmgLOxO6zGOGadxF6NmNAhcGp9Lg,9024
 warpzone/deltastorage/schema.py,sha256=AhkJnsB9oX3HVwoCUWcYSdD1zwouT2wj5XydXgU5WEA,4483
 warpzone/deltastorage/slicing.py,sha256=r4Fbg8HHBkZlrYEC847yq86caCTTE02JzpnY_TqI_MA,638
-warpzone/deltastorage/store.py,sha256=bmWfIy5TXQU1G89-IuBrK4xZ5s7UtDz18QVuJ-ZECBw,4337
-warpzone/deltastorage/table.py,sha256=JE5rviTlmKICRG6u0dFxnLBzkCfT7y0I77pOkxXIXO0,5212
+warpzone/deltastorage/store.py,sha256=LwoDXk0xtQuzfMbqv_ldIRh-9RzA-Hx9NRhtBmlbWB8,4812
+warpzone/deltastorage/table.py,sha256=_5tBQEKnimNWsAWyVhGjuX_3pllqzakkMrvALw6-2Qw,5713
 warpzone/enums/__init__.py,sha256=Gvd-EeZc5jWu5WdusYe3i3Zyt1v85CljBK_EJu7EwZk,24
 warpzone/enums/topicenum.py,sha256=cj9F7kYzwvDeHN5jJXwgnHXTh6E64VwRNpFepNthIyk,213
 warpzone/function/__init__.py,sha256=rJOZBpWsUgjMc7YtXMJ1rLGm45KB1AhDJ_Y2ISiSISc,35
 warpzone/function/checks.py,sha256=B9YqThymf16ac_fVAYKilv20ru5v9nwXgHlbxYIaG98,1018
 warpzone/function/functionize.py,sha256=bSV0QvwKbD9Vo3a_8cc1rgV2rzTdMMvidinyXItBfvs,2128
-warpzone/function/integrations.py,sha256=sDt2BTx6a4mVc-33wTITP9XQVPustwj7rkX4urTyOqo,4018
+warpzone/function/integrations.py,sha256=Law-0TI_tbm8rq5XXuilDH69_0LhoqaZhDbTL498Qik,4016
 warpzone/function/monitor.py,sha256=xD13d4795a9qgGphOywFl4sOXRQjKypiL2ozQFOseqQ,1862
 warpzone/function/process.py,sha256=nbUVywM8ChfUwuaqFisgaD98aNRgeZkK4g5sbtuBdRs,2339
 warpzone/function/processors/__init__.py,sha256=DhIdSWLBcIeSO8IJdxPqGIhgwwnkDN6_Xqwy93BCLeA,46
@@ -28,8 +28,8 @@ warpzone/function/types.py,sha256=5m2hRrnLC3eqIlAH5-MM9_wKjMZ6lYawZtCOVStyFuY,72
 warpzone/healthchecks/__init__.py,sha256=9gc_Mt2szs8sDSwy0V4l3JZ6d9hX41xTpZCkDP2qsY4,2108
 warpzone/healthchecks/model.py,sha256=mM7DnrirLbUpBPPfi82MUPP654D0eOR2_F65TmzsPD0,1187
 warpzone/monitor/__init__.py,sha256=gXT2cxR4tlZER54zd7D49ZQBVyitLaqj13_cUoILuyM,109
-warpzone/monitor/logs.py,sha256=AO853uEpD9ZZldQsBic5IR9k3GuQPyWafk4pC5fAGio,3084
-warpzone/monitor/traces.py,sha256=2ANxKmO1yca3JFBi68P9HeXQo8WVOOVeT1ev3GLURSY,4107
+warpzone/monitor/logs.py,sha256=q3SUQCtG1ii0B9GkVs2l8kgZ5b5bI2qI6L0SoOW2QTY,3147
+warpzone/monitor/traces.py,sha256=Xc_po1LxJFy5jtNWxIVphIInl_d89Zw3Rb21PsdQhQA,4170
 warpzone/servicebus/data/__init__.py,sha256=lnc0uiaGLF0qMi_rWhCpRSFvaj0CJEiMCAl6Yqn1ZiA,21
 warpzone/servicebus/data/client.py,sha256=zECS3JwedhYnDk8PntYgIYpBF_uu9YN38KzpPFK7CKs,6511
 warpzone/servicebus/events/__init__.py,sha256=lnc0uiaGLF0qMi_rWhCpRSFvaj0CJEiMCAl6Yqn1ZiA,21
@@ -37,7 +37,7 @@ warpzone/servicebus/events/client.py,sha256=8v8XsF-2RwzKIi_93IzR_eR-BZTGXXHSuV4P
 warpzone/servicebus/events/triggers.py,sha256=_QuPTBbje7LrBoz0qhhgrtDZOcE6x1S9GNu-WJUQ8bY,2626
 warpzone/tablestorage/db/__init__.py,sha256=lnc0uiaGLF0qMi_rWhCpRSFvaj0CJEiMCAl6Yqn1ZiA,21
 warpzone/tablestorage/db/base_client.py,sha256=ropKO6z0UXqBl38NuGYV4VZ_ZFm4w1d84ReOLYoBKLY,2376
-warpzone/tablestorage/db/client.py,sha256=5I_5pUzB9ZVYyj8cf2g9PIuMdXIfaytWsAnDm6EZnyg,6661
+warpzone/tablestorage/db/client.py,sha256=-OchZI7x--Z9msNp912ggZ_sNGwNDXO-i7VGV6INSrg,6587
 warpzone/tablestorage/db/table_config.py,sha256=PC45dnr3vVMVr8ktu3GWoH8u3JKJaCRAvMCroUaC1NE,1336
 warpzone/tablestorage/tables/__init__.py,sha256=l_8wElG1oam39fmXqIXDAnCYqbKXEyX1G0cD6JCCv4s,78
 warpzone/tablestorage/tables/client.py,sha256=chvmTz5S8M1-IWYgfJHyTCwLZVsVzlchMuCYwHzQhCE,3636
@@ -52,6 +52,6 @@ warpzone/tools/copy.py,sha256=5fddotMZkXZO8avzUbGOhvs0cp8mce95pNpy0oPVjnQ,2596
 warpzone/transform/__init__.py,sha256=ruGa7tl-v4ndlWpULE1jSGU_a4_iRc3V6eyNr5xKP9E,27
 warpzone/transform/data.py,sha256=Abb8PcrgMbbNCJkkIUdtrTHdlY0OfXid387qw1nDpFY,2362
 warpzone/transform/schema.py,sha256=nbSQtDMvXkyqGKuwhuFCF0WsEDsaNyoPYpMKvbsKlv8,2423
-warpzone_sdk-15.0.0.dev14.dist-info/METADATA,sha256=C7hvF9-5FlHq-EwoLv65npYpGLZTIEXAsRXo_0Sp_GE,7285
-warpzone_sdk-15.0.0.dev14.dist-info/WHEEL,sha256=3ny-bZhpXrU6vSQ1UPG34FoxZBp3lVcvK0LkgUz6VLk,88
-warpzone_sdk-15.0.0.dev14.dist-info/RECORD,,
+warpzone_sdk-15.1.0.dist-info/METADATA,sha256=q0pdO2gAFw7uT3fTYGgBfLeFDOyQstIvEQR-3qfKZb8,7279
+warpzone_sdk-15.1.0.dist-info/WHEEL,sha256=3ny-bZhpXrU6vSQ1UPG34FoxZBp3lVcvK0LkgUz6VLk,88
+warpzone_sdk-15.1.0.dist-info/RECORD,,