cloe-nessy 0.3.16.6b0__py3-none-any.whl → 0.3.17.0__py3-none-any.whl

cloe_nessy/integration/delta_loader/__init__.py

@@ -0,0 +1,14 @@
+from .delta_loader import DeltaLoader
+from .delta_loader_factory import DeltaLoaderFactory, DeltaLoadOptions, consume_delta_load
+from .strategies import DeltaCDFConfig, DeltaCDFLoader, DeltaTimestampConfig, DeltaTimestampLoader
+
+__all__ = [
+    "consume_delta_load",
+    "DeltaCDFConfig",
+    "DeltaCDFLoader",
+    "DeltaLoader",
+    "DeltaLoaderFactory",
+    "DeltaLoadOptions",
+    "DeltaTimestampConfig",
+    "DeltaTimestampLoader",
+]

cloe_nessy/integration/delta_loader/delta_load_options.py

@@ -0,0 +1,37 @@
+from pathlib import Path
+from typing import Self
+
+import yaml
+from pydantic import BaseModel
+
+
+class DeltaLoadOptions(BaseModel):
+    """Options to configure the DeltaLoader.
+
+    Args:
+        strategy: Delta load strategy to use.
+        delta_load_identifier: Unique delta load identifier used to track the delta load metadata.
+        strategy_options: Options used to configure the chosen delta load strategy.
+            See the config class of the particular strategy for more info.
+        metadata_table_identifier: Identifier of the metadata table used to keep
+            track of the delta load metadata. The table will be created if it does
+            not exist. If none, it will default to `<source_catalog>.<source_schema>.metadata_delta_load`.
+    """
+
+    strategy: str
+    delta_load_identifier: str
+    strategy_options: dict
+    metadata_table_identifier: str | None = None
+
+    @classmethod
+    def from_yaml_str(cls, yaml_str: str) -> Self:
+        """Creates an instance of DeltaLoadOptions from a YAML string."""
+        options = yaml.safe_load(yaml_str)
+        return cls(**options)
+
+    @classmethod
+    def from_file(cls, path: str | Path) -> Self:
+        """Creates an instance of DeltaLoadOptions from a YAML file."""
+        with Path(path).open() as f:
+            yaml_str = f.read()
+        return cls.from_yaml_str(yaml_str)

cloe_nessy/integration/delta_loader/delta_loader.py

@@ -0,0 +1,165 @@
+from abc import ABC, abstractmethod
+from datetime import datetime
+from functools import partial
+
+from delta import DeltaTable  # type: ignore[import-untyped]
+from pyspark.sql import DataFrame
+from pyspark.sql import functions as F
+
+from ...integration.writer import CatalogWriter
+from ...logging import LoggerMixin
+from ...object_manager import TableManager
+from ...session import SessionManager
+from .delta_loader_metadata_table import DeltaLoaderMetadataTable
+
+
+class DeltaLoader(ABC, LoggerMixin):
+    """Base class for delta load operations.
+
+    Args:
+        table_identifier: Identifier for the table to be loaded.
+        delta_load_identifier: Identifier for the delta load.
+        metadata_table_identifier: Identifier for the metadata table. If None,
+            the metadata_table_identifier will be derived from the table identifier:
+            `<table_catalog>.<table_schema>.metadata_delta_load`.
+    """
+
+    def __init__(
+        self,
+        table_identifier: str,
+        delta_load_identifier: str,
+        metadata_table_identifier: str | None = None,
+    ):
+        self._spark = SessionManager.get_spark_session()
+        self._console_logger = self.get_console_logger()
+        self.table_identifier = table_identifier
+        self.delta_load_identifier = delta_load_identifier
+        self.metadata_table_identifier = (
+            metadata_table_identifier
+            or f"{self.table_identifier.split('.')[0]}.{self.table_identifier.split('.')[1]}.metadata_delta_load"
+        )
+        table_manager = TableManager()
+        table_manager.create_table(table=DeltaLoaderMetadataTable(identifier=self.metadata_table_identifier))
+
+    @abstractmethod
+    def read_data(
+        self,
+        options: dict[str, str] | None = None,
+    ) -> DataFrame:
+        """Reads data incrementally using a strategy.
+
+        Args:
+            options: Additional DataFrameReader options.
+        """
+        pass
+
+    @abstractmethod
+    def verify(self) -> None:
+        """Verify that the source table qualifies for the delta load strategy."""
+        pass
+
+    def _query(self, query: str) -> DataFrame:
+        df = self._spark.sql(query)
+        return df
+
+    def _create_metadata_entry(
+        self,
+        *,
+        rows: int,
+        last_read_timestamp: datetime | None = None,
+        start_version: int | None = None,
+        end_version: int | None = None,
+        start_commit_timestamp: datetime | None = None,
+        end_commit_timestamp: datetime | None = None,
+    ) -> None:
+        """Creates an entry in the metadata table for the delta load."""
+        self._console_logger.info(
+            f"Creating metadata entry for table: [ {self.table_identifier} ] with Delta Load Identifier: [ {self.delta_load_identifier} ]",
+        )
+        metadata_df = self._spark.range(1)
+        metadata_df = metadata_df.select(
+            F.lit(rows).alias("rows").cast("bigint"),
+            F.lit(False).alias("is_processed"),
+            F.lit(False).alias("is_stale"),
+            F.lit(self.table_identifier).alias("source_table_identifier"),
+            F.lit(self.delta_load_identifier).alias("delta_load_identifier"),
+            F.lit(start_version).alias("start_commit_version"),
+            F.lit(end_version).alias("end_commit_version"),
+            F.lit(start_commit_timestamp).alias("start_commit_timestamp_utc"),
+            F.lit(end_commit_timestamp).alias("end_commit_timestamp_utc"),
+            F.lit(last_read_timestamp).alias("last_read_timestamp"),
+            F.current_timestamp().alias("__DCR"),
+            F.current_timestamp().alias("__DMR"),
+        ).withColumn(
+            "BK",
+            F.md5(
+                F.concat_ws(
+                    "-",
+                    F.col("source_table_identifier"),
+                    F.col("delta_load_identifier"),
+                    F.current_timestamp(),
+                ),
+            ),
+        )
+        catalog_writer = CatalogWriter()
+        catalog_writer.write_table(
+            df=metadata_df,
+            table_identifier=self.metadata_table_identifier,
+            mode="append",
+        )
+
+    def _invalidate_versions(self) -> None:
+        """Invalidate any pending changes in the metadata for the delta load."""
+        self._console_logger.info(
+            f"Invalidating unprocessed delta load metadata for table: [ {self.table_identifier} ] with Delta Load Identifier: [ {self.delta_load_identifier} ]",
+        )
+        delta_table = DeltaTable.forName(self._spark, self.metadata_table_identifier)
+        delta_table.update(
+            condition=(F.col("source_table_identifier") == self.table_identifier)
+            & (F.col("delta_load_identifier") == self.delta_load_identifier)
+            & ~F.col("is_processed")
+            & ~F.col("is_stale"),
+            set={"is_stale": F.lit(True), "__DMR": F.current_timestamp()},
+        )
+
+    def reset_cdf(self) -> None:
+        """Invalidates all changes in the metadata for the delta load."""
+        delta_table = DeltaTable.forName(self._spark, self.metadata_table_identifier)
+        self._console_logger.info(
+            f"Resetting delta load metadata for table: [ {self.table_identifier} ] with Delta Load Identifier: [ {self.delta_load_identifier} ]",
+        )
+        delta_table.update(
+            condition=(F.col("source_table_identifier") == self.table_identifier)
+            & (F.col("delta_load_identifier") == self.delta_load_identifier)
+            & ~F.col("is_stale"),
+            set={"is_stale": F.lit(True), "__DMR": F.current_timestamp()},
+        )
+
+    def consume_data(self) -> None:
+        """Marks data as consumed in the metadata for the delta load."""
+        df = self._spark.table(self.metadata_table_identifier)
+        df = df.filter(
+            (F.col("source_table_identifier") == self.table_identifier)
+            & (F.col("delta_load_identifier") == self.delta_load_identifier)
+            & ~F.col("is_processed")
+            & ~F.col("is_stale"),
+        )
+        df = df.groupBy("BK", "delta_load_identifier").agg(F.max("__DCR")).limit(1)
+        self._console_logger.info(
+            f"Mark metadata for table as processed: [ {self.table_identifier} ] with Delta Load Identifier: [ {self.delta_load_identifier} ].",
+        )
+        delta_table = DeltaTable.forName(self._spark, self.metadata_table_identifier)
+        delta_table.alias("target").merge(df.alias("source"), "target.BK = source.BK").whenMatchedUpdate(
+            set={
+                "is_processed": F.lit(True),
+                "__DMR": F.current_timestamp(),
+            },
+        ).execute()
+
+    def write_data(self, write_callable: partial):
+        """Wrapper to write and consume a delta load."""
+        try:
+            write_callable()
+        except Exception as e:
+            raise RuntimeError("Error while writing...") from e
+        self.consume_data()

cloe_nessy/integration/delta_loader/delta_loader_factory.py

@@ -0,0 +1,53 @@
+from typing import Any
+
+from .delta_load_options import DeltaLoadOptions
+from .delta_loader import DeltaLoader
+from .strategies import DeltaCDFConfig, DeltaCDFLoader, DeltaTimestampConfig, DeltaTimestampLoader
+
+
+def consume_delta_load(
+    runtime_info: dict[str, Any],
+    delta_load_identifier: str | None = None,
+) -> None:
+    """Consumes a delta load by updating the metadata table.
+
+    Args:
+        runtime_info: Runtime information.
+        delta_load_identifier: If set, the ConsumeDeltaLoadAction action
+            will only consume DeltaLoader transaction for the given
+            delta_load_identifier.
+    """
+    for table_name, value in runtime_info["delta_load_options"].items():
+        if delta_load_identifier is None or delta_load_identifier == value.get("delta_load_identifier"):
+            delta_loader: DeltaLoader = DeltaLoaderFactory.create_loader(
+                table_identifier=table_name,
+                options=DeltaLoadOptions(
+                    **value,
+                ),
+            )
+            delta_loader.consume_data()
+
+
+class DeltaLoaderFactory:
+    """Factory to create a DeltaLoader instance based on the DeltaLoadOptions."""
+
+    @staticmethod
+    def create_loader(table_identifier: str, options: DeltaLoadOptions) -> DeltaLoader:
+        """Creates an instance of DeltaLoader, choosing the desired strategy."""
+        if options.strategy.upper() == "CDF":
+            cdf_config = DeltaCDFConfig(**options.strategy_options)
+            return DeltaCDFLoader(
+                table_identifier=table_identifier,
+                delta_load_identifier=options.delta_load_identifier,
+                config=cdf_config,
+                metadata_table_identifier=options.metadata_table_identifier,
+            )
+        if options.strategy.upper() == "TIMESTAMP":
+            timestamp_config = DeltaTimestampConfig(**options.strategy_options)
+            return DeltaTimestampLoader(
+                table_identifier=table_identifier,
+                delta_load_identifier=options.delta_load_identifier,
+                config=timestamp_config,
+                metadata_table_identifier=options.metadata_table_identifier,
+            )
+        raise ValueError(f"Unknown strategy: {options.strategy}")

cloe_nessy/integration/delta_loader/delta_loader_metadata_table.py

@@ -0,0 +1,68 @@
+from ...models import Column, Table
+
+
+class DeltaLoaderMetadataTable(Table):
+    """A Table Model for the Delta CDF Reader Metadata Table."""
+
+    data_source_format: str = "DELTA"
+    partition_by: list[str] = ["source_table_identifier"]
+    liquid_clustering: bool = True
+    columns: list[Column] = [
+        Column(
+            name="BK",
+            data_type="STRING",
+        ),
+        Column(
+            name="delta_load_identifier",
+            data_type="STRING",
+        ),
+        Column(
+            name="source_table_identifier",
+            data_type="STRING",
+        ),
+        Column(
+            name="is_processed",
+            data_type="BOOLEAN",
+        ),
+        Column(
+            name="is_stale",
+            data_type="BOOLEAN",
+        ),
+        Column(
+            name="last_read_timestamp",
+            data_type="TIMESTAMP",
+            nullable=True,
+        ),
+        Column(
+            name="rows",
+            data_type="BIGINT",
+        ),
+        Column(
+            name="start_commit_version",
+            data_type="INT",
+            nullable=True,
+        ),
+        Column(
+            name="end_commit_version",
+            data_type="INT",
+            nullable=True,
+        ),
+        Column(
+            name="start_commit_timestamp_utc",
+            data_type="TIMESTAMP",
+            nullable=True,
+        ),
+        Column(
+            name="end_commit_timestamp_utc",
+            data_type="TIMESTAMP",
+            nullable=True,
+        ),
+        Column(
+            name="__DCR",
+            data_type="TIMESTAMP",
+        ),
+        Column(
+            name="__DMR",
+            data_type="TIMESTAMP",
+        ),
+    ]

cloe_nessy/integration/delta_loader/strategies/__init__.py

@@ -0,0 +1,9 @@
+from .delta_cdf_loader import DeltaCDFConfig, DeltaCDFLoader
+from .delta_timestamp_loader import DeltaTimestampConfig, DeltaTimestampLoader
+
+__all__ = [
+    "DeltaCDFConfig",
+    "DeltaCDFLoader",
+    "DeltaTimestampConfig",
+    "DeltaTimestampLoader",
+]