fabricks 3.0.11__py3-none-any.whl

fabricks/cdc/base/configurator.py

@@ -0,0 +1,223 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import List, Optional, Union
+
+from pyspark.sql import DataFrame, SparkSession
+
+from fabricks.cdc.base._types import AllowedSources
+from fabricks.context import SPARK
+from fabricks.context.log import DEFAULT_LOGGER
+from fabricks.metastore.database import Database
+from fabricks.metastore.table import Table
+from fabricks.utils._types import DataFrameLike
+
+
+class Configurator(ABC):
+    def __init__(
+        self,
+        database: str,
+        *levels: str,
+        change_data_capture: str,
+        spark: Optional[SparkSession] = None,
+    ):
+        if spark is None:
+            spark = SPARK
+        assert spark is not None
+        self.spark: SparkSession = spark
+
+        self.database = Database(database)
+        self.levels = levels
+        self.change_data_capture = change_data_capture
+        self.table = Table(self.database.name, *self.levels, spark=self.spark)
+
+    @property
+    def is_view(self):
+        return self.table.is_view
+
+    @property
+    def registered(self):
+        return self.table.registered
+
+    @property
+    def qualified_name(self):
+        return f"{self.database}_{'_'.join(self.levels)}"
+
+    @abstractmethod
+    def get_query(self, src: AllowedSources, **kwargs) -> str: ...
+
+    @abstractmethod
+    def get_data(self, src: AllowedSources, **kwargs) -> DataFrame: ...
+
+    @abstractmethod
+    def create_table(
+        self,
+        src: AllowedSources,
+        partitioning: Optional[bool] = False,
+        partition_by: Optional[Union[List[str], str]] = None,
+        identity: Optional[bool] = False,
+        liquid_clustering: Optional[bool] = False,
+        cluster_by: Optional[Union[List[str], str]] = None,
+        properties: Optional[dict[str, str]] = None,
+        **kwargs,
+    ): ...
+
+    @abstractmethod
+    def drop(self): ...
+
+    @abstractmethod
+    def create_or_replace_view(self, src: Union[Table, str], **kwargs): ...
+
+    @property
+    def allowed_input__columns(self) -> List[str]:
+        cols = self.__columns
+
+        if self.slowly_changing_dimension:
+            if "__valid_from" in cols:
+                cols.remove("__valid_from")
+            if "__valid_to" in cols:
+                cols.remove("__valid_to")
+            if "__is_current" in cols:
+                cols.remove("__is_current")
+            if "__is_deleted" in cols:
+                cols.remove("__is_deleted")
+
+        return cols
+
+    @property
+    def allowed_ouput_leading__columns(self) -> List[str]:
+        cols = [
+            "__identity",
+            "__source",
+            "__key",
+            "__timestamp",
+            "__valid_from",
+            "__valid_to",
+            "__is_current",
+            "__is_deleted",
+        ]
+
+        if self.change_data_capture == "scd1":
+            cols.remove("__valid_from")
+            cols.remove("__valid_to")
+        elif self.change_data_capture == "scd2":
+            cols.remove("__timestamp")
+
+        return cols
+
+    @property
+    def allowed_output_trailing__columns(self) -> List[str]:
+        cols = [
+            "__operation",
+            "__metadata",
+            "__hash",
+            "__rescued_data",
+        ]
+
+        if self.slowly_changing_dimension:
+            cols.remove("__operation")
+
+        return cols
+
+    @property
+    def __columns(self) -> List[str]:
+        return [
+            # Leading
+            "__identity",
+            "__source",
+            "__key",
+            "__timestamp",
+            "__valid_from",
+            "__valid_to",
+            "__is_current",
+            "__is_deleted",
+            # Trailing
+            "__operation",
+            "__metadata",
+            "__hash",
+            "__rescued_data",
+        ]
+
+    @property
+    def slowly_changing_dimension(self) -> bool:
+        return self.change_data_capture in ["scd1", "scd2"]
+
+    def get_src(self, src: AllowedSources) -> DataFrame:
+        if isinstance(src, DataFrameLike):
+            df = src
+        elif isinstance(src, Table):
+            df = self.table.dataframe
+        elif isinstance(src, str):
+            df = self.spark.sql(src)
+        else:
+            raise ValueError(f"{src} not allowed")
+
+        return df
+
+    def has_data(self, src: AllowedSources, **kwargs) -> bool:
+        DEFAULT_LOGGER.debug("check if has data", extra={"label": self})
+        df = self.get_src(src=src)
+        return not df.isEmpty()
+
+    def get_columns(
+        self,
+        src: AllowedSources,
+        backtick: Optional[bool] = True,
+        sort: Optional[bool] = True,
+        check: Optional[bool] = True,
+    ) -> List[str]:
+        if backtick:
+            backtick = True
+
+        df = self.get_src(src=src)
+        columns = df.columns
+
+        if check:
+            for c in columns:
+                # avoid duplicate column issue in merge
+                if c.startswith("__") and c in self.__columns:
+                    assert c in self.allowed_input__columns, f"{c} is not allowed"
+
+        if sort:
+            columns = self.sort_columns(columns)
+
+        if backtick:
+            return [f"`{c}`" for c in columns]
+        else:
+            return columns
+
+    def sort_columns(self, columns: List[str]) -> List[str]:
+        fields = [c for c in columns if not c.startswith("__")]
+
+        leading = self.allowed_ouput_leading__columns
+        trailing = self.allowed_output_trailing__columns
+
+        # move __hash to the front of the table to ensure statistics are present
+        if "__key" not in columns and "__hash" in columns:
+            leading = ["__hash" if c == "__key" else c for c in leading]
+            trailing = [c for c in trailing if c != "__hash"]
+
+        __leading = [c for c in leading if c in columns]
+        __trailing = [c for c in trailing if c in columns]
+
+        return __leading + fields + __trailing
+
+    def reorder_dataframe(self, df: DataFrame) -> DataFrame:
+        columns = self.sort_columns(df.columns)
+        columns = [f"`{c}`" for c in columns]
+        return df.select(columns)
+
+    @abstractmethod
+    def optimize_table(self): ...
+
+    @abstractmethod
+    def update_schema(self, src: AllowedSources, **kwargs): ...
+
+    @abstractmethod
+    def get_differences_with_deltatable(self, src: AllowedSources, **kwargs): ...
+
+    @abstractmethod
+    def overwrite_schema(self, src: AllowedSources): ...
+
+    def __str__(self):
+        return f"{self.table.qualified_name}"

fabricks/cdc/base/generator.py

@@ -0,0 +1,177 @@
+from __future__ import annotations
+
+from typing import Any, List, Optional, Sequence, Union, cast
+
+from py4j.protocol import Py4JJavaError
+from pyspark.sql import DataFrame
+
+from fabricks.cdc.base._types import AllowedSources
+from fabricks.cdc.base.configurator import Configurator
+from fabricks.context.log import DEFAULT_LOGGER
+from fabricks.metastore.table import SchemaDiff, Table
+from fabricks.utils._types import DataFrameLike
+from fabricks.utils.sqlglot import fix as fix_sql
+
+
+class Generator(Configurator):
+    def drop(self):
+        self.table.drop()
+
+    def create_table(
+        self,
+        src: AllowedSources,
+        partitioning: Optional[bool] = False,
+        partition_by: Optional[Union[List[str], str]] = None,
+        identity: Optional[bool] = False,
+        liquid_clustering: Optional[bool] = False,
+        cluster_by: Optional[Union[List[str], str]] = None,
+        properties: Optional[dict[str, str]] = None,
+        masks: Optional[dict[str, str]] = None,
+        primary_key: Optional[dict[str, Any]] = None,
+        foreign_keys: Optional[dict[str, Any]] = None,
+        comments: Optional[dict[str, str]] = None,
+        **kwargs,
+    ):
+        kwargs["mode"] = "complete"
+        kwargs["slice"] = False
+        kwargs["rectify"] = False
+        kwargs["deduplicate"] = False
+
+        df = self.get_data(src, **kwargs)
+
+        if partitioning is True:
+            assert partition_by, "partitioning column(s) not found"
+
+        df = self.reorder_dataframe(df)
+
+        identity = False if identity is None else identity
+        liquid_clustering = False if liquid_clustering is None else liquid_clustering
+
+        self.table.create(
+            df=df,
+            partitioning=partitioning,
+            partition_by=partition_by,
+            identity=identity,
+            liquid_clustering=liquid_clustering,
+            cluster_by=cluster_by,
+            properties=properties,
+            masks=masks,
+            primary_key=primary_key,
+            foreign_keys=foreign_keys,
+            comments=comments,
+        )
+
+    def create_or_replace_view(self, src: Union[Table, str], schema_evolution: bool = True, **kwargs):
+        assert not isinstance(src, DataFrameLike), "dataframe not allowed"
+
+        assert kwargs["mode"] == "complete", f"{kwargs['mode']} not allowed"
+        sql = self.get_query(src, **kwargs)
+
+        df = self.spark.sql(sql)
+        df = self.reorder_dataframe(df)
+        columns = [f"`{c}`" for c in df.columns]
+
+        sql = f"""
+        create or replace view {self}
+        {"with schema evolution" if schema_evolution else "-- no schema evolution"}
+        as
+        with __view as (
+          {sql}
+        )
+        select
+          {",".join(columns)}
+        from __view
+        """
+        sql = fix_sql(sql)
+        DEFAULT_LOGGER.debug("create or replace view", extra={"label": self, "sql": sql})
+
+        try:
+            self.spark.sql(sql)
+        except Py4JJavaError as e:
+            DEFAULT_LOGGER.exception("fail to execute sql query", extra={"label": self, "sql": sql}, exc_info=e)
+
+    def optimize_table(self):
+        columns = None
+
+        if self.change_data_capture == "scd1":
+            columns = ["__key"]
+        elif self.change_data_capture == "scd2":
+            columns = ["__key", "__valid_from"]
+
+        self.table.optimize(columns=columns)
+
+    def get_differences_with_deltatable(self, src: AllowedSources, **kwargs) -> DataFrame:
+        from pyspark.sql.types import StringType, StructField, StructType
+
+        schema = StructType(
+            [
+                StructField("column", StringType(), False),
+                StructField("data_type", StringType(), True),
+                StructField("new_column", StringType(), True),
+                StructField("new_data_type", StringType(), True),
+                StructField("status", StringType(), True),
+            ]
+        )
+
+        if self.is_view:
+            return self.spark.createDataFrame([], schema=schema)
+
+        else:
+            kwargs["mode"] = "complete"
+            if "slice" in kwargs:
+                del kwargs["slice"]
+
+            df = self.get_data(src, **kwargs)
+            df = self.reorder_dataframe(df)
+
+            diffs = self.table.get_schema_differences(df)
+            return self.spark.createDataFrame([cast(Any, d.model_dump()) for d in diffs], schema=schema)
+
+    def get_schema_differences(self, src: AllowedSources, **kwargs) -> Optional[Sequence[SchemaDiff]]:
+        if self.is_view:
+            return None
+
+        else:
+            kwargs["mode"] = "complete"
+            if "slice" in kwargs:
+                del kwargs["slice"]
+
+            df = self.get_data(src, **kwargs)
+            df = self.reorder_dataframe(df)
+
+            return self.table.get_schema_differences(df)
+
+    def schema_drifted(self, src: AllowedSources, **kwargs) -> Optional[bool]:
+        d = self.get_schema_differences(src, **kwargs)
+        if d is None:
+            return None
+        return len(d) > 0
+
+    def _update_schema(
+        self,
+        src: AllowedSources,
+        overwrite: bool = False,
+        widen_types: bool = False,
+        **kwargs,
+    ):
+        if self.is_view:
+            assert not isinstance(src, DataFrameLike), "dataframe not allowed"
+            self.create_or_replace_view(src=src)
+
+        else:
+            kwargs["mode"] = "complete"
+            if "slice" in kwargs:
+                del kwargs["slice"]
+
+            df = self.get_data(src, **kwargs)
+            df = self.reorder_dataframe(df)
+            if overwrite:
+                self.table.overwrite_schema(df)
+            else:
+                self.table.update_schema(df, widen_types=widen_types)
+
+    def update_schema(self, src: AllowedSources, **kwargs):
+        self._update_schema(src=src, **kwargs)
+
+    def overwrite_schema(self, src: AllowedSources, **kwargs):
+        self._update_schema(src=src, overwrite=True, **kwargs)

fabricks/cdc/base/merger.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+from typing import Optional, Union
+
+from jinja2 import Environment, PackageLoader
+from pyspark.sql import DataFrame
+
+from fabricks.cdc.base._types import AllowedSources
+from fabricks.cdc.base.processor import Processor
+from fabricks.context.log import DEFAULT_LOGGER
+from fabricks.metastore.view import create_or_replace_global_temp_view
+from fabricks.utils._types import DataFrameLike
+from fabricks.utils.sqlglot import fix as fix_sql
+
+
+class Merger(Processor):
+    def get_merge_context(self, src: Union[DataFrame, str], **kwargs) -> dict:
+        if isinstance(src, DataFrameLike):
+            format = "dataframe"
+            columns = self.get_columns(src, backtick=False, sort=False, check=False)  # already done in processor
+        elif isinstance(src, str):
+            format = "view"
+            columns = self.get_columns(
+                f"select * from {src}", backtick=False, sort=False, check=False
+            )  # already done in processor
+        else:
+            raise ValueError(f"{src} not allowed")
+
+        assert "__merge_key" in columns, "__merge_key not found"
+        assert "__merge_condition" in columns, "__merge_condition not found"
+
+        keys = kwargs.get("keys")
+        if isinstance(keys, str):
+            keys = [keys]
+
+        columns = [c for c in columns if c not in ["__merge_condition", "__merge_key"]]
+        fields = [c for c in columns if not c.startswith("__")]
+        where = kwargs.get("update_where") if self.table.rows > 0 else None
+        soft_delete = "__is_deleted" in columns
+
+        has_source = "__source" in columns
+        has_key = "__key" in columns
+        has_metadata = "__metadata" in columns
+        has_hash = "__hash" in columns
+        has_timestamp = "__timestamp" in columns
+        has_identity = "__identity" in columns
+
+        # 'NoneType' object is not iterable
+        if keys:
+            keys = [f"`{k}`" for k in keys]
+        if columns:
+            columns = [f"`{c}`" for c in columns]
+        if fields:
+            fields = [f"`{c}`" for c in fields]
+
+        assert "__key" or keys, f"{self} - __key or keys not found"
+
+        return {
+            "src": src,
+            "format": format,
+            "tgt": self.table,
+            "cdc": self.change_data_capture,
+            "columns": columns,
+            "fields": fields,
+            "soft_delete": soft_delete,
+            "has_source": has_source,
+            "has_identity": has_identity,
+            "has_key": has_key,
+            "has_hash": has_hash,
+            "keys": keys,
+            "has_metadata": has_metadata,
+            "has_timestamp": has_timestamp,
+            "where": where,
+        }
+
+    def get_merge_query(self, src: Union[DataFrame, str], fix: Optional[bool] = True, **kwargs) -> str:
+        context = self.get_merge_context(src=src, **kwargs)
+        environment = Environment(loader=PackageLoader("fabricks.cdc", "templates"))
+        merge = environment.get_template("merge.sql.jinja")
+
+        try:
+            sql = merge.render(**context)
+        except Exception as e:
+            DEFAULT_LOGGER.debug("context", extra={"label": self, "content": context})
+            raise e
+
+        if fix:
+            try:
+                sql = sql.replace("{src}", "src")
+                sql = fix_sql(sql)
+                sql = sql.replace("`src`", "{src}")
+                DEFAULT_LOGGER.debug("merge", extra={"label": self, "sql": sql})
+
+            except Exception as e:
+                DEFAULT_LOGGER.exception("fail to clean sql query", extra={"label": self, "sql": sql})
+                raise e
+
+        return sql
+
+    def merge(self, src: AllowedSources, **kwargs):
+        if not self.table.exists():
+            self.create_table(src, **kwargs)
+
+        df = self.get_data(src, **kwargs)
+        global_temp_view = f"{self.qualified_name}__merge"
+        view = create_or_replace_global_temp_view(global_temp_view, df, uuid=kwargs.get("uuid", False), job=self)
+
+        merge = self.get_merge_query(view, **kwargs)
+        DEFAULT_LOGGER.debug("exec merge", extra={"label": self, "sql": merge})
+        self.spark.sql(merge, src=view)