fabricks 2024.7.1.5__py3-none-any.whl

fabricks/cdc/base/configurator.py

@@ -0,0 +1,145 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import List, Optional, Union
+
+from databricks.sdk.runtime import dbutils as _dbutils
+from databricks.sdk.runtime import spark as _spark
+from pyspark.sql import DataFrame, SparkSession
+
+from fabricks.metastore.database import Database
+from fabricks.metastore.table import Table
+
+
+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.dbutils = _dbutils
+
+        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)
+
+    def is_view(self):
+        return self.table.is_view()
+
+    def registered(self):
+        return self.table.registered()
+
+    @abstractmethod
+    def get_query(self, src: Union[DataFrame, Table, str], **kwargs):
+        raise NotImplementedError()
+
+    @abstractmethod
+    def get_data(self, src: Union[DataFrame, Table, str], **kwargs) -> DataFrame:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def create_table(
+        self,
+        src: Union[DataFrame, Table, str],
+        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,
+    ):
+        raise NotImplementedError()
+
+    @abstractmethod
+    def drop(self):
+        raise NotImplementedError()
+
+    @abstractmethod
+    def create_or_replace_view(self, query: Optional[str] = None):
+        raise NotImplementedError()
+
+    @property
+    def allowed_leading_columns(self):
+        cols = ["__identity", "__key", "__timestamp", "__valid_from", "__valid_to"]
+        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_trailing_columns(self):
+        cols = [
+            "__source",
+            "__operation",
+            "__is_current",
+            "__is_deleted",
+            "__metadata",
+            "__hash",
+            "__rescued_data",
+        ]
+        if self.change_data_capture == "scd1":
+            cols.remove("__operation")
+        elif self.change_data_capture == "scd2":
+            cols.remove("__operation")
+        return cols
+
+    @property
+    def slowly_changing_dimension(self) -> bool:
+        return self.change_data_capture in ["scd1", "scd2"]
+
+    def get_src(self, src: Union[DataFrame, Table, str]) -> DataFrame:
+        if isinstance(src, DataFrame):
+            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 get_columns(self, src: Union[DataFrame, Table, str], backtick: Optional[bool] = True) -> List[str]:
+        if backtick:
+            backtick = True
+
+        df = self.get_src(src=src)
+        columns = df.columns
+
+        if backtick:
+            return [f"`{c}`" for c in columns]
+        else:
+            return columns
+
+    def reorder_columns(self, df: DataFrame) -> DataFrame:
+        fields = [f"`{c}`" for c in df.columns if not c.startswith("__")]
+        __leading = [c for c in self.allowed_leading_columns if c in df.columns]
+        __trailing = [c for c in self.allowed_trailing_columns if c in df.columns]
+
+        columns = __leading + fields + __trailing
+
+        return df.select(columns)
+
+    @abstractmethod
+    def optimize_table(self):
+        raise NotImplementedError()
+
+    @abstractmethod
+    def update_schema(self, **kwargs):
+        raise NotImplementedError()
+
+    @abstractmethod
+    def overwrite_schema(self):
+        raise NotImplementedError()
+
+    def __str__(self):
+        return f"{self.table.qualified_name}"

fabricks/cdc/base/generator.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+from typing import List, Optional, Union
+
+from py4j.protocol import Py4JJavaError
+from pyspark.sql import DataFrame
+
+from fabricks.cdc.base.configurator import Configurator
+from fabricks.context.log import Logger
+from fabricks.metastore.table import Table
+from fabricks.utils.sqlglot import fix as fix_sql
+
+
+class Generator(Configurator):
+    def drop(self):
+        self.table.drop()
+
+    def create_table(
+        self,
+        src: Union[DataFrame, Table, str],
+        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,
+    ):
+        kwargs["mode"] = "complete"
+        kwargs["filter"] = False
+        kwargs["rectify"] = False
+        kwargs["deduplicate"] = False
+        df = self.get_data(src, **kwargs)
+
+        if liquid_clustering:
+            assert cluster_by, "clustering column not found"
+        elif partitioning:
+            assert partition_by, "partitioning column not found"
+
+        fields = [c for c in df.columns if not c.startswith("__")]
+        __leading = [c for c in self.allowed_leading_columns if c in df.columns]
+        __trailing = [c for c in self.allowed_trailing_columns if c in df.columns]
+        columns = __leading + fields + __trailing
+
+        df = df.select([f"`{c}`" for c in columns])
+
+        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,
+        )
+
+    def create_or_replace_view(self, src: Union[Table, str], **kwargs):
+        assert not isinstance(src, DataFrame), "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_columns(df)
+        columns = [f"`{c}`" for c in df.columns]
+
+        sql = f"""
+        create or replace view {self}
+        as
+        with __view as (
+          {sql}
+        )
+        select
+          {','.join(columns)}
+        from __view
+        """
+        sql = fix_sql(sql)
+        Logger.debug("create or replace view", extra={"job": self, "sql": sql})
+
+        try:
+            self.spark.sql(sql)
+        except Py4JJavaError:
+            Logger.exception("🙈", extra={"job": self})
+
+    def optimize_table(self):
+        liquid_clustering = self.table.get_property("delta.feature.liquid") == "supported"
+        if liquid_clustering:
+            self.table.optimize()
+        else:
+            columns = None
+            if self.change_data_capture == "scd1":
+                columns = ["__key"]
+            elif self.change_data_capture == "scd2":
+                columns = ["__key", "__valid_from"]
+            vorder = self.table.get_property("delta.parquet.vorder.enabled") or "false"
+            vorder = vorder.lower() == "true"
+            self.table.optimize(columns=columns, vorder=vorder)
+
+    def update_schema(self, src: Union[DataFrame, Table, str], **kwargs):
+        overwrite = kwargs.get("overwrite", False)
+
+        if self.is_view():
+            assert not isinstance(src, DataFrame), "dataframe not allowed"
+            self.create_or_replace_view(src=src, **kwargs)
+        else:
+            kwargs["mode"] = "complete"
+            df = self.get_data(src, **kwargs)
+            if overwrite:
+                self.table.overwrite_schema(df)
+            else:
+                self.table.update_schema(df)
+
+    def overwrite_schema(self, src: Union[DataFrame, Table, str]):
+        self.update_schema(src=src, overwrite=True)

fabricks/cdc/base/merger.py

@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+from typing import Optional, Union
+
+from jinja2 import Environment, PackageLoader
+from pyspark.sql import DataFrame
+
+from fabricks.cdc.base.processor import Processor
+from fabricks.context.log import Logger
+from fabricks.metastore.table import Table
+from fabricks.metastore.view import create_or_replace_global_temp_view
+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, DataFrame):
+            format = "dataframe"
+            columns = self.get_columns(src, backtick=False)
+        elif isinstance(src, str):
+            format = "view"
+            columns = self.get_columns(f"select * from {src}", backtick=False)
+        else:
+            raise ValueError(f"{src} not allowed")
+
+        assert "__merge_key" in columns
+        assert "__merge_condition" in columns
+
+        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:
+            Logger.debug("context", extra={"job": self, "content": context})
+            raise e
+
+        if fix:
+            try:
+                sql = sql.replace("{src}", "src")
+                sql = fix_sql(sql)
+                sql = sql.replace("`src`", "{src}")
+                Logger.debug("merge", extra={"job": self, "sql": sql})
+            except Exception as e:
+                Logger.exception("🙈", extra={"job": self, "sql": sql})
+                raise e
+        else:
+            Logger.debug("merge", extra={"job": self, "sql": sql})
+
+        return sql
+
+    def merge(self, src: Union[DataFrame, Table, str], **kwargs):
+        if not self.table.exists():
+            self.create_table(src, **kwargs)
+
+        df = self.get_data(src, **kwargs)
+        if df:
+            global_temp_view = f"{self.database}_{'_'.join(self.levels)}__merge"
+            view = create_or_replace_global_temp_view(global_temp_view, df, uuid=kwargs.get("uuid", False))
+
+            merge = self.get_merge_query(view, **kwargs)
+            self.spark.sql(merge, src=view)