fabricks 2024.7.1.5__py3-none-any.whl

fabricks/metastore/relational.py

@@ -0,0 +1,61 @@
+from typing import Optional
+
+from databricks.sdk.runtime import dbutils as _dbutils
+from databricks.sdk.runtime import spark as _spark
+from pyspark.errors.exceptions.base import AnalysisException
+from pyspark.sql import SparkSession
+
+from fabricks.context.log import Logger
+from fabricks.metastore.database import Database
+
+
+class Relational:
+    def __init__(self, database: str, *levels: str, spark: Optional[SparkSession] = None):
+        self.database = Database(database)
+        self.levels = levels
+        if spark is None:
+            spark = _spark
+        assert spark is not None
+        self.spark: SparkSession = spark
+        self.dbutils = _dbutils
+
+    @property
+    def name(self) -> str:
+        return "_".join(self.levels)
+
+    @property
+    def qualified_name(self) -> str:
+        return f"{self.database.name}.{self.name}"
+
+    def registered(self):
+        try:
+            df = self.spark.sql(f"show tables in {self.database}").where(f"tableName == '{self.name}'")
+            return not df.isEmpty()
+        # not found
+        except AnalysisException:
+            return False
+
+    def is_view(self):
+        try:
+            df = self.spark.sql(f"show views in {self.database}").where(f"viewName == '{self.name}'")
+            return not df.isEmpty()
+        # not found
+        except AnalysisException:
+            return False
+
+    def is_table(self):
+        if self.is_view():
+            return False
+        else:
+            return self.registered()
+
+    def drop(self):
+        if self.is_view():
+            Logger.warning("drop view from metastore", extra={"job": self})
+            self.spark.sql(f"drop view if exists {self}")
+        elif self.is_table():
+            Logger.warning("drop table from metastore", extra={"job": self})
+            self.spark.sql(f"drop table if exists {self}")
+
+    def __str__(self):
+        return self.qualified_name

fabricks/metastore/table.py

@@ -0,0 +1,529 @@
+import re
+from typing import List, Optional, Union, overload
+
+from databricks.sdk.runtime import spark
+from delta import DeltaTable
+from pyspark.errors.exceptions.base import AnalysisException
+from pyspark.sql import DataFrame, SparkSession
+from pyspark.sql.functions import expr, max
+from pyspark.sql.types import StructType
+from typing_extensions import deprecated
+
+from fabricks.context.log import Logger
+from fabricks.metastore.relational import Relational
+from fabricks.utils.path import Path
+from fabricks.utils.sqlglot import fix
+
+
+class Table(Relational):
+    @classmethod
+    def from_step_topic_item(cls, step: str, topic: str, item: str, spark: Optional[SparkSession] = spark):
+        return cls(step, topic, item, spark=spark)
+
+    @property
+    @deprecated("use delta_path instead")
+    def deltapath(self) -> Path:
+        return self.database.deltapath.join("/".join(self.levels))
+
+    @property
+    def delta_path(self) -> Path:
+        return self.database.deltapath.join("/".join(self.levels))
+
+    @property
+    def deltatable(self) -> DeltaTable:
+        return DeltaTable.forPath(self.spark, self.deltapath.string)
+
+    @property
+    def delta_table(self) -> DeltaTable:
+        return DeltaTable.forPath(self.spark, self.deltapath.string)
+
+    @property
+    def dataframe(self) -> DataFrame:
+        return self.spark.sql(f"select * from {self}")
+
+    @property
+    def columns(self) -> List[str]:
+        return self.dataframe.columns
+
+    @property
+    def rows(self) -> int:
+        return self.spark.sql(f"select count(*) from {self}").collect()[0][0]
+
+    @property
+    def last_version(self) -> int:
+        df = self.describe_history()
+        version = df.select(max("version")).collect()[0][0]
+        return version
+
+    def drop(self):
+        super().drop()
+        if self.deltapath.exists():
+            Logger.debug("delete delta folder", extra={"job": self})
+            self.deltapath.rm()
+
+    @overload
+    def create(
+        self,
+        df: DataFrame,
+        *,
+        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,
+    ): ...
+
+    @overload
+    def create(
+        self,
+        *,
+        schema: StructType,
+        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,
+    ): ...
+
+    def create(
+        self,
+        df: Optional[DataFrame] = None,
+        schema: Optional[StructType] = None,
+        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,
+    ):
+        self._create(
+            df=df,
+            schema=schema,
+            partitioning=partitioning,
+            partition_by=partition_by,
+            identity=identity,
+            liquid_clustering=liquid_clustering,
+            cluster_by=cluster_by,
+            properties=properties,
+        )
+
+    def _create(
+        self,
+        df: Optional[DataFrame] = None,
+        schema: Optional[StructType] = None,
+        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,
+    ):
+        Logger.info("create table", extra={"job": self})
+        if not df:
+            assert schema is not None
+            df = self.spark.createDataFrame([], schema)
+
+        def _backtick(name: str, dtype: str) -> str:
+            j = df.schema[name].jsonValue()
+            r = re.compile(r"(?<='name': ')[^']+(?=',)")
+            names = re.findall(r, str(j))
+            for n in names:
+                escaped = re.escape(n)
+                dtype = re.sub(f"(?<=,){escaped}(?=:)|(?<=<){escaped}(?=:)", f"`{n}`", dtype)
+            return dtype
+
+        ddl_columns = ",\n\t".join([f"`{name}` {_backtick(name, dtype)}" for name, dtype in df.dtypes])
+        ddl_identity = "-- no identity" if "__identity" not in df.columns else ""
+        ddl_cluster_by = "-- no cluster by"
+        ddl_partition_by = "-- no partitioned by"
+        ddl_tblproperties = "-- not tblproperties"
+
+        if liquid_clustering:
+            assert cluster_by
+            if isinstance(cluster_by, str):
+                cluster_by = [cluster_by]
+            cluster_by = [f"`{c}`" for c in cluster_by]
+            ddl_cluster_by = "cluster by (" + ", ".join(cluster_by) + ")"
+        if partitioning:
+            assert partition_by
+            if isinstance(partition_by, str):
+                partition_by = [partition_by]
+            partition_by = [f"`{p}`" for p in partition_by]
+            ddl_partition_by = "partitioned by (" + ", ".join(partition_by) + ")"
+
+        if identity:
+            ddl_identity = "__identity bigint generated by default as identity (start with 1 increment by 1), "
+
+        if not properties:
+            special_char = False
+            for c in df.columns:
+                match = re.search(r"[^a-zA-Z0-9_]", c)
+                if match:
+                    special_char = True
+                    break
+            if special_char:
+                properties = {
+                    "delta.columnMapping.mode": "name",
+                    "delta.minReaderVersion": "2",
+                    "delta.minWriterVersion": "5",
+                }
+        if properties:
+            ddl_tblproperties = (
+                "tblproperties (" + ",".join(f"'{key}' = '{value}'" for key, value in properties.items()) + ")"
+            )
+
+        sql = f""" 
+        create table if not exists {self.qualified_name} 
+        (
+        {ddl_identity}
+        {ddl_columns}
+        )
+        {ddl_tblproperties}
+        {ddl_partition_by}
+        {ddl_cluster_by}
+        location '{self.deltapath}'
+        """
+        try:
+            sql = fix(sql)
+        except Exception:
+            pass
+        Logger.debug("ddl", extra={"job": self, "sql": sql})
+        self.spark.sql(sql)
+
+    def is_deltatable(self) -> bool:
+        return DeltaTable.isDeltaTable(self.spark, str(self.deltapath))
+
+    def column_mapping_enabled(self) -> bool:
+        return self.get_property("delta.columnMapping.mode") == "name"
+
+    def exists(self) -> bool:
+        return self.is_deltatable() and self.registered()
+
+    def register(self):
+        Logger.debug("register table", extra={"job": self})
+        self.spark.sql(f"create table if not exists {self.qualified_name} using delta location '{self.deltapath}'")
+
+    def restore_to_version(self, version: int):
+        Logger.info(f"restore table to version {version}", extra={"job": self})
+        self.spark.sql(f"restore table {self.qualified_name} to version as of {version}")
+
+    def truncate(self):
+        Logger.warning("truncate table", extra={"job": self})
+        self.create_restore_point()
+        self.spark.sql(f"truncate table {self.qualified_name}")
+
+    def schema_drifted(self, df: DataFrame) -> bool:
+        return not self._check_schema_drift(df).isEmpty()
+
+    def _check_schema_drift(self, df: DataFrame) -> DataFrame:
+        Logger.debug("check schema drift", extra={"job": self})
+
+        new_df = self.spark.createDataFrame(df.dtypes, ["new_name", "new_type"])  # type: ignore
+        new_df = new_df.filter(~new_df.new_name.startswith("__"))
+
+        old_df = self.spark.createDataFrame(self.dataframe.dtypes, ["old_name", "old_type"])  # type: ignore
+        old_df = old_df.filter(~old_df.old_name.startswith("__"))
+
+        cond = [new_df["new_name"] == old_df["old_name"]]
+        df_diff = (
+            new_df.join(old_df, on=cond, how="outer")
+            .where(
+                """
+                coalesce(old_name, -1) <> coalesce(new_name, -1) 
+                or 
+                coalesce(old_type, -1) <> coalesce(new_type, -1)
+                """
+            )
+            .withColumn(
+                "operation",
+                expr("if(new_name is null, 'drop', if(old_name is null, 'add', 'update'))"),
+            )
+            .withColumn("column", expr("coalesce(new_name, old_name)"))
+        )
+        return df_diff
+
+    def _fix_schema(self, df: DataFrame, overwrite: bool = False):
+        drift_df = self._check_schema_drift(df)
+
+        if not drift_df.isEmpty():
+            Logger.info("update table", extra={"job": self})
+            todo_df = drift_df.where("operation in ('add', 'update')")
+            if not todo_df.isEmpty():
+                for row in todo_df.collect():
+                    if row.operation == "add":
+                        Logger.debug(f"add column { row.column}", extra={"job": self})
+                    else:
+                        Logger.debug(
+                            "update column { row.column} ({row.old_type} -> {row.new_type})",
+                            extra={"job": self},
+                        )
+
+                    try:
+                        col_df = df.select(row.column).where("1 == 2")
+                        (
+                            self.deltatable.alias("dt")
+                            .merge(col_df.alias("df"), "1 == 2")
+                            .whenNotMatchedInsertAll()
+                            .execute()
+                        )
+                    except Exception:
+                        pass
+
+            if overwrite:
+                drift_df = self._check_schema_drift(df)
+                Logger.warning("overwrite table", extra={"job": self})
+                for row in drift_df.collect():
+                    if row.operation == "add":
+                        self.add_column(row.column, row.new_type)
+                    elif row.operation == "drop":
+                        self.drop_column(row.column)
+                    elif row.operation == "update":
+                        try:
+                            self.change_column(row.column, row.new_type)
+                        except AnalysisException:
+                            self.drop_column(row.column)
+                            self.add_column(row.column, row.new_type)
+                    else:
+                        raise ValueError(f"{row.operation} not allowed")
+
+    def update_schema(self, df: DataFrame):
+        self._fix_schema(df, overwrite=False)
+
+    def overwrite_schema(self, df: DataFrame):
+        self._fix_schema(df, overwrite=True)
+
+    def vacuum(self, retention_days: int = 7):
+        Logger.debug(f"vacuum table (removing files older than {retention_days} days)", extra={"job": self})
+        self.spark.sql("SET self.spark.databricks.delta.retentionDurationCheck.enabled = False")
+        try:
+            self.create_restore_point()
+            retention_hours = retention_days * 24
+            self.deltatable.vacuum(retention_hours)
+        finally:
+            # finally
+            pass
+        self.spark.sql("SET self.spark.databricks.delta.retentionDurationCheck.enabled = True")
+
+    def optimize(
+        self,
+        columns: Optional[Union[str, List[str]]] = None,
+        vorder: Optional[bool] = False,
+    ):
+        Logger.info("optimize", extra={"job": self})
+
+        zorder_by = columns is not None
+        if zorder_by:
+            if isinstance(columns, str):
+                columns = [columns]
+            columns = [f"`{c}`" for c in columns]
+            cols = ", ".join(columns)
+
+            if vorder:
+                Logger.debug(f"zorder by {cols} vorder", extra={"job": self})
+                self.spark.sql(f"optimize {self.qualified_name} zorder by ({cols}) vorder")
+            else:
+                Logger.debug(f"zorder by {cols}", extra={"job": self})
+                self.spark.sql(f"optimize {self.qualified_name} zorder by ({cols})")
+
+        elif vorder:
+            Logger.debug("vorder", extra={"job": self})
+            self.spark.sql(f"optimize {self.qualified_name} vorder")
+
+        else:
+            Logger.debug("optimize", extra={"job": self})
+            self.spark.sql(f"optimize {self.qualified_name}")
+
+    def analyze(self):
+        Logger.debug("analyze", extra={"job": self})
+        self.compute_statistics()
+        self.compute_delta_statistics()
+
+    def compute_statistics(self):
+        Logger.debug("compute statistics", extra={"job": self})
+        cols = [
+            f"`{name}`"
+            for name, dtype in self.dataframe.dtypes
+            if not dtype.startswith("struct") and not dtype.startswith("array") and name not in ["__metadata"]
+        ]
+        cols = ", ".join(sorted(cols))
+        self.spark.sql(f"analyze table delta.`{self.deltapath}` compute statistics for columns {cols}")
+
+    def compute_delta_statistics(self):
+        Logger.debug("compute delta statistics", extra={"job": self})
+        self.spark.sql(f"analyze table delta.`{self.deltapath}` compute delta statistics")
+
+    def drop_column(self, name: str):
+        assert self.column_mapping_enabled(), "column mapping not enabled"
+        Logger.warning(f"drop column {name}", extra={"job": self})
+        self.spark.sql(
+            f"""
+            alter table {self.qualified_name}
+            drop column `{name}`
+            """
+        )
+
+    def change_column(self, name: str, type: str):
+        assert self.column_mapping_enabled(), "column mapping not enabled"
+        Logger.info(f"change column {name} ({type})", extra={"job": self})
+        self.spark.sql(
+            f"""
+            alter table {self.qualified_name}
+            change column `{name}` `{name}` {type}
+            """
+        )
+
+    def rename_column(self, old: str, new: str):
+        assert self.column_mapping_enabled(), "column mapping not enabled"
+        Logger.info(f"rename column {old} -> {new}", extra={"job": self})
+        self.spark.sql(
+            f"""
+            alter table {self.qualified_name}
+            rename column `{old}` to `{new}`
+            """
+        )
+
+    def get_details(self) -> DataFrame:
+        return self.spark.sql(f"describe detail {self.qualified_name}")
+
+    def get_properties(self) -> DataFrame:
+        return self.spark.sql(f"show tblproperties {self.qualified_name}")
+
+    def get_description(self) -> DataFrame:
+        return self.spark.sql(f"describe extended {self.qualified_name}")
+
+    def get_history(self) -> DataFrame:
+        df = self.spark.sql(f"describe history {self.qualified_name}")
+        return df
+
+    def get_last_version(self) -> int:
+        df = self.get_history()
+        version = df.select(max("version")).collect()[0][0]
+        return version
+
+    def get_property(self, key: str) -> Optional[str]:
+        try:
+            value = self.get_properties().where(f"key == '{key}'").select("value").collect()[0][0]
+            return value
+
+        except IndexError:
+            return None
+
+    def enable_change_data_feed(self):
+        Logger.debug("enable change data feed", extra={"job": self})
+        self.set_property("delta.enableChangeDataFeed", "true")
+
+    def enable_column_mapping(self):
+        Logger.debug("enable column mapping", extra={"job": self})
+        try:
+            self.spark.sql(
+                f"""
+                alter table {self.qualified_name} 
+                set tblproperties ('delta.columnMapping.mode' = 'name')
+                """
+            )
+        except Exception:
+            Logger.debug("update reader and writer version", extra={"job": self})
+            self.spark.sql(
+                f"""
+                alter table {self.qualified_name} 
+                set tblproperties (
+                    'delta.columnMapping.mode' = 'name',
+                    'delta.minReaderVersion' = '2',
+                    'delta.minWriterVersion' = '5'
+                )
+                """
+            )
+
+    def set_property(self, key: Union[str, int], value: Union[str, int]):
+        Logger.debug(f"set property {key} = {value}", extra={"job": self})
+        self.spark.sql(
+            f"""
+            alter table {self.qualified_name}
+            set tblproperties ({key} = '{value}')
+            """
+        )
+
+    def add_constraint(self, name: str, expr: str):
+        Logger.debug(f"add constraint ({name} check ({expr}))", extra={"job": self})
+        self.spark.sql(
+            f"""
+            alter table {self.qualified_name}
+            add constraint {name} check ({expr});
+            """
+        )
+
+    def add_comment(self, comment: str):
+        Logger.debug(f"add comment '{comment}'", extra={"job": self})
+        self.spark.sql(
+            f"""
+            comment on table {self.qualified_name}
+            is '{comment}';
+            """
+        )
+
+    def add_materialized_column(self, name: str, expr: str, type: str):
+        assert self.column_mapping_enabled(), "column mapping not enabled"
+        Logger.info(f"add materialized column ({name} {type})", extra={"job": self})
+        self.spark.sql(
+            f""""
+            alter table {self.qualified_name}
+            add columns (`{name}` {type} materialized {expr})
+            """
+        )
+
+    def add_column(self, name: str, type: str, after: Optional[str] = None):
+        Logger.info(f"add column {name} ({type})", extra={"job": self})
+        ddl_after = "" if not after else f"after {after}"
+        self.spark.sql(
+            f"""
+            alter table {self.qualified_name}
+            add columns (`{name}` {type} {ddl_after})
+            """
+        )
+
+    def create_bloomfilter_index(self, columns: Union[str, List[str]]):
+        if isinstance(columns, str):
+            columns = [columns]
+        columns = [f"`{c}`" for c in columns]
+        cols = ", ".join(columns)
+
+        Logger.info(f"bloomfilter by {cols}", extra={"job": self})
+        self.spark.sql(
+            f"""
+            create bloomfilter index on table {self.qualified_name}
+            for columns ({cols})
+            """
+        )
+
+    def create_restore_point(self):
+        last_version = self.get_last_version() + 1
+        self.set_property("fabricks.last_version", last_version)
+
+    def show_properties(self) -> DataFrame:
+        return self.spark.sql(f"show tblproperties {self.qualified_name}")
+
+    def describe_detail(self) -> DataFrame:
+        return self.spark.sql(f"describe detail {self.qualified_name}")
+
+    def describe_extended(self) -> DataFrame:
+        return self.spark.sql(f"describe extended {self.qualified_name}")
+
+    def describe_history(self) -> DataFrame:
+        df = self.spark.sql(f"describe history {self.qualified_name}")
+        return df
+
+    def enable_liquid_clustering(self, columns: Union[str, List[str]]):
+        if isinstance(columns, str):
+            columns = [columns]
+        columns = [f"`{c}`" for c in columns]
+        cols = ", ".join(columns)
+        Logger.info(f"cluster by {cols}", extra={"job": self})
+
+        self.spark.sql(
+            f"""
+            alter table {self.qualified_name}
+            cluster by ({cols})
+            """
+        )

fabricks/metastore/utils.py

@@ -0,0 +1,35 @@
+from databricks.sdk.runtime import spark
+from pyspark.sql import DataFrame
+
+
+def get_tables(schema: str) -> DataFrame:
+    table_df = spark.sql(f"show tables in {schema}")
+    view_df = spark.sql(f"show views in {schema}")
+    df = spark.sql(
+        """
+        select 
+            database,
+            concat_ws('.', database, tableName) as table
+        from 
+            {tables} 
+            left anti join {views} on tableName = viewName
+        """,
+        tables=table_df,
+        views=view_df,
+    )
+    return df
+
+
+def get_views(schema: str) -> DataFrame:
+    view_df = spark.sql(f"show views in {schema}")
+    df = spark.sql(
+        """
+        select 
+            namespace as database,
+            concat_ws('.', namespace, viewName) as view
+        from 
+            {views}
+        """,
+        views=view_df,
+    )
+    return df

fabricks/metastore/view.py

@@ -0,0 +1,40 @@
+from typing import Optional, Union
+from uuid import uuid4
+
+import pandas as pd
+from databricks.sdk.runtime import spark as _spark
+from pyspark.sql import DataFrame, SparkSession
+
+from fabricks.context.log import Logger
+from fabricks.metastore.relational import Relational
+
+
+class View(Relational):
+    @staticmethod
+    def create_or_replace(
+        df: Union[DataFrame, pd.DataFrame],
+        *dependencies,
+        spark: Optional[SparkSession] = None,
+    ) -> str:
+        if spark is None:
+            spark = _spark
+        assert spark is not None
+
+        uuid = str(uuid4().hex)
+        df = spark.createDataFrame(df) if isinstance(df, pd.DataFrame) else df
+        if dependencies:
+            for d in dependencies:
+                df = df.join(d.where("1 == 2"), how="leftanti")
+
+        df.createOrReplaceGlobalTempView(uuid)
+        return uuid
+
+
+def create_or_replace_global_temp_view(name: str, df: DataFrame, uuid: Optional[bool] = False) -> str:
+    if uuid:
+        name = f"{name}__{str(uuid4().hex)}"
+
+    job = name.split("__")[0]
+    Logger.debug("create global temp view", extra={"job": job})
+    df.createOrReplaceGlobalTempView(name)
+    return f"global_temp.{name}"

fabricks/utils/README.md

@@ -0,0 +1,3 @@
+# BMS DNA Fabricks Databricks
+
+Databricks - Fabricks