fabricks 2024.7.1.5__py3-none-any.whl

fabricks/core/jobs/gold.py

@@ -0,0 +1,218 @@
+import re
+from typing import List, Optional, cast
+
+from databricks.sdk.runtime import dbutils
+from pyspark.sql import DataFrame
+
+from fabricks.cdc.nocdc import NoCDC
+from fabricks.context.log import Logger
+from fabricks.core.jobs.base.job import BaseJob
+from fabricks.core.jobs.base.types import TGold
+from fabricks.core.udfs import is_registered, register_udf
+from fabricks.metastore.view import create_or_replace_global_temp_view
+from fabricks.utils.path import Path
+
+
+class Gold(BaseJob):
+    def __init__(
+        self, step: TGold, topic: Optional[str] = None, item: Optional[str] = None, job_id: Optional[str] = None
+    ):  # type: ignore
+        super().__init__(
+            "gold",
+            step=step,
+            topic=topic,
+            item=item,
+            job_id=job_id,
+        )
+
+    _sql: Optional[str] = None
+    _sql_path: Optional[Path] = None
+    _schema_drift: Optional[bool] = None
+
+    @classmethod
+    def from_job_id(cls, step: str, job_id: str):
+        return cls(step=cast(TGold, step), job_id=job_id)
+
+    @classmethod
+    def from_step_topic_item(cls, step: str, topic: str, item: str):
+        return cls(step=cast(TGold, step), topic=topic, item=item)
+
+    @property
+    def stream(self) -> bool:
+        return False
+
+    @property
+    def schema_drift(self) -> bool:
+        if not self._schema_drift:
+            _schema_drift = self.step_conf.get("options", {}).get("schema_drift", False)
+            assert _schema_drift is not None
+            self._schema_drift = cast(bool, _schema_drift)
+        return self._schema_drift
+
+    @property
+    def persist(self) -> bool:
+        return self.mode in ["update", "append", "complete"]
+
+    @property
+    def virtual(self) -> bool:
+        return self.mode in ["memory"]
+
+    def get_sql(self) -> str:
+        return self.paths.runtime.get_sql()
+
+    def get_udfs(self) -> List[str]:
+        # udf not allowed in invoke
+        if self.mode == "invoke":
+            return []
+        # udf not allowed in notebook
+        elif self.options.job.get("notebook"):
+            return []
+        # udf not allowed in table
+        elif self.options.job.get("table"):
+            return []
+        else:
+            matches = []
+            if "udf_" in self.get_sql():
+                r = re.compile(r"(?<=udf_)\w*(?=\()")
+                matches = re.findall(r, self.get_sql())
+                matches = set(matches)
+                matches = list(matches)
+            return matches
+
+    def register_udfs(self):
+        for u in self.get_udfs():
+            if not is_registered(u):
+                Logger.debug(f"register udf ({u})", extra={"job": self})
+                register_udf(udf=u, spark=self.spark)
+
+    def base_transform(self, df: DataFrame) -> DataFrame:
+        df = df.transform(self.extender)
+        return df
+
+    def get_data(self, stream=False, transform: Optional[bool] = False) -> DataFrame:
+        if self.options.job.get_boolean("requirements"):
+            import sys
+
+            sys.path.append("/dbfs/mnt/fabricks/site-packages")
+
+        if self.mode == "invoke":
+            df = self.spark.createDataFrame([{}])  # type: ignore
+
+        elif self.options.job.get("notebook"):
+            Logger.debug("run notebook", extra={"job": self})
+            path = self.paths.runtime.get_notebook_path()
+            global_temp_view = dbutils.notebook.run(path, self.timeouts.job, arguments={})
+            df = self.spark.sql(f"select * from global_temp.{global_temp_view}")
+
+        elif self.options.job.get("table"):
+            table = self.options.job.get("table")
+            df = self.spark.read.table(table)  # type: ignore
+
+        else:
+            assert self.get_sql(), "sql not found"
+            self.register_udfs()
+            df = self.spark.sql(self.get_sql())
+
+        if transform:
+            df = self.base_transform(df)
+        return df
+
+    def create_or_replace_view(self):
+        assert self.mode == "memory", f"{self.mode} not allowed"
+
+        df = self.spark.sql(self.get_sql())
+        cdc_options = self.get_cdc_context(df)
+        self.cdc.create_or_replace_view(self.get_sql(), **cdc_options)
+
+    def get_cdc_context(self, df: DataFrame) -> dict:
+        if "__order_duplicate_by_asc" in df.columns:
+            order_duplicate_by = {"__order_duplicate_by_asc": "asc"}
+        elif "__order_duplicate_by_desc" in df.columns:
+            order_duplicate_by = {"__order_duplicate_by_desc": "desc"}
+        else:
+            order_duplicate_by = None
+
+        context = {
+            "add_metadata": True,
+            "soft_delete": True if self.slowly_changing_dimension else None,
+            "deduplicate_key": self.options.job.get_boolean("deduplicate", None),
+            "deduplicate_hash": True if self.slowly_changing_dimension else None,
+            "deduplicate": False,  # assume no duplicate in gold
+            "rectify": False,  # assume no reload in gold
+            "order_duplicate_by": order_duplicate_by,
+        }
+
+        if self.slowly_changing_dimension:
+            if "__key" not in df.columns:
+                context["add_key"] = True
+            if "__hash" not in df.columns:
+                context["add_hash"] = True
+
+            if "__operation" not in df.columns:
+                context["deduplicate_hash"] = None  # assume no duplicate hash
+                if self.mode == "update":
+                    context["add_operation"] = "reload"
+                    context["rectify"] = True
+                else:
+                    context["add_operation"] = "upsert"
+
+        if self.mode == "update" and self.change_data_capture == "scd2":
+            context["filter"] = "update"
+
+        if self.mode == "memory":
+            context["mode"] = "complete"
+
+        return context
+
+    def for_each_batch(self, df: DataFrame, batch: Optional[int] = None):
+        assert self.persist, f"{self.mode} not allowed"
+
+        context = self.get_cdc_context(df=df)
+
+        # if dataframe, reference is passed (BUG)
+        name = f"{self.step}_{self.topic}_{self.item}"
+        global_temp_view = create_or_replace_global_temp_view(name=name, df=df)
+        sql = f"select * from {global_temp_view}"
+
+        if self.mode == "update":
+            assert not isinstance(self.cdc, NoCDC), "nocdc update not allowed"
+            self.cdc.update(sql, **context)
+        elif self.mode == "append":
+            assert isinstance(self.cdc, NoCDC), f"{self.change_data_capture} append not allowed"
+            self.cdc.append(sql, **context)
+        elif self.mode == "complete":
+            self.cdc.complete(sql, **context)
+        else:
+            raise ValueError(f"{self.mode} - not allowed")
+
+        self.check_duplicate_key()
+        self.check_duplicate_hash()
+
+    def for_each_run(self, schedule: Optional[str] = None):
+        if self.mode == "invoke":
+            self.invoke(schedule=schedule)
+        else:
+            super().for_each_run(schedule=schedule)
+
+    def create(self):
+        if self.mode == "invoke":
+            Logger.info("invoke (no table nor view)", extra={"job": self})
+        else:
+            super().create()
+
+    def register(self):
+        if self.mode == "invoke":
+            Logger.info("invoke (no table nor view)", extra={"job": self})
+        else:
+            super().register()
+
+    def optimize(
+        self,
+        vacuum: Optional[bool] = True,
+        optimize: Optional[bool] = True,
+        analyze: Optional[bool] = True,
+    ):
+        if self.mode == "memory":
+            Logger.debug("memory (no optimize)", extra={"job": self})
+        else:
+            super().optimize()

fabricks/core/jobs/silver.py

@@ -0,0 +1,354 @@
+from typing import Optional, cast
+
+from pyspark.sql import DataFrame, Row
+from pyspark.sql.functions import expr
+
+from fabricks.cdc.nocdc import NoCDC
+from fabricks.context.log import Logger
+from fabricks.core.jobs.base.job import BaseJob
+from fabricks.core.jobs.base.types import TBronze, TSilver
+from fabricks.core.jobs.bronze import Bronze
+from fabricks.metastore.view import create_or_replace_global_temp_view
+from fabricks.utils.helpers import concat_dfs
+from fabricks.utils.read.read import read
+from fabricks.utils.sqlglot import fix as fix_sql
+
+
+class Silver(BaseJob):
+    def __init__(
+        self, step: TSilver, topic: Optional[str] = None, item: Optional[str] = None, job_id: Optional[str] = None
+    ):  # type: ignore
+        super().__init__(
+            "silver",
+            step=step,
+            topic=topic,
+            item=item,
+            job_id=job_id,
+        )
+
+    _parent_step: Optional[TBronze] = None
+    _stream: Optional[bool] = None
+
+    @classmethod
+    def from_job_id(cls, step: str, job_id: str):
+        return cls(step=cast(TSilver, step), job_id=job_id)
+
+    @classmethod
+    def from_step_topic_item(cls, step: str, topic: str, item: str):
+        return cls(step=cast(TSilver, step), topic=topic, item=item)
+
+    @property
+    def stream(self) -> bool:
+        if not self._stream:
+            _stream = self.options.job.get("stream")
+            if _stream is None:
+                _stream = self.step_conf.get("options", {}).get("stream")
+            self._stream = _stream if _stream is not None else True
+        return self._stream  # type: ignore
+
+    @property
+    def schema_drift(self) -> bool:
+        return True
+
+    @property
+    def persist(self) -> bool:
+        return self.mode in ["update", "append", "latest"]
+
+    @property
+    def virtual(self) -> bool:
+        return self.mode in ["combine", "memory"]
+
+    @property
+    def parent_step(self) -> TBronze:
+        if not self._parent_step:
+            _parent_step = self.step_conf.get("options", {}).get("parent")
+            _parent_step = cast(TBronze, _parent_step)
+            assert _parent_step is not None
+            self._parent_step = _parent_step
+        return self._parent_step
+
+    def base_transform(self, df: DataFrame) -> DataFrame:
+        df = df.transform(self.extender)
+        if "__metadata" in df.columns:
+            df = df.withColumn(
+                "__metadata",
+                expr(
+                    """
+                    struct(
+                        __metadata.file_path as file_path,
+                        __metadata.file_name as file_name,
+                        __metadata.file_size as file_size,            
+                        __metadata.file_modification_time as file_modification_time,
+                        __metadata.inserted as inserted,
+                    cast(current_date() as timestamp) as updated
+                    )
+                    """
+                ),
+            )
+        return df
+
+    def get_data(self, stream: bool = True, transform: Optional[bool] = False) -> DataFrame:
+        dep_df = self.get_dependencies()
+        assert dep_df, "not dependency found"
+        dep_df = dep_df.orderBy("parent_id")
+        dependencies = dep_df.count()
+
+        if self.mode == "memory":
+            assert dependencies == 1, f"more than 1 dependency not allowed ({dependencies})"
+
+            parent = dep_df.collect()[0].parent
+            df = self.spark.sql(f"select * from {parent}")
+
+        elif self.mode == "combine":
+            dfs = []
+            for row in dep_df.collect():
+                df = self.spark.sql(f"select * from {row.parent}")
+                dfs.append(df)
+            df = concat_dfs(dfs)
+
+        else:
+            dfs = []
+
+            for row in dep_df.collect():
+                try:
+                    bronze = Bronze.from_job_id(step=self.parent_step, job_id=row["parent_id"])
+                    if bronze.mode in ["memory", "register"]:
+                        # data already transformed if bronze is persisted
+                        df = bronze.get_data(stream=stream, transform=True)
+                    else:
+                        df = read(
+                            stream=stream,
+                            path=bronze.table.deltapath,
+                            file_format="delta",
+                            metadata=False,
+                            spark=self.spark,
+                        )
+                    if dependencies > 1:
+                        assert "__source" in df.columns, "__source not found"
+                    dfs.append(df)
+                except Exception as e:
+                    Logger.exception("🙈", extra={"job": self})
+                    raise e
+
+            df = concat_dfs(dfs)
+
+        # transforms
+        df = self.filter_where(df)
+        df = self.encrypt(df)
+        if transform:
+            df = self.base_transform(df)
+        return df
+
+    def get_dependencies(self, df: Optional[DataFrame] = None) -> Optional[DataFrame]:
+        dependencies = []
+        parents = self.options.job.get_list("parents") or []
+        if parents:
+            for p in parents:
+                dependencies.append(Row(self.job_id, p, "job"))
+        else:
+            p = f"{self.parent_step}.{self.topic}_{self.item}"
+            dependencies.append(Row(self.job_id, p, "parser"))
+
+        if dependencies:
+            Logger.debug(f"dependencies ({', '.join([row[1] for row in dependencies])})", extra={"job": self})
+            df = self.spark.createDataFrame(dependencies, schema=["job_id", "parent", "origin"])
+            df = df.transform(self.add_dependency_details)
+            return df
+
+    def create_or_replace_view(self):
+        assert self.mode in ["memory", "combine"], f"{self.mode} not allowed"
+
+        dep_df = self.get_dependencies()
+        assert dep_df, "dependency not found"
+
+        if self.mode == "combine":
+            queries = []
+
+            for row in dep_df.collect():
+                columns = self.get_data().columns
+                df = self.spark.sql(f"select * from {row.parent}")
+                cols = [f"`{c}`" if c in df.columns else f"null as `{c}`" for c in columns if c not in ["__source"]]
+                source = "__source" if "__source" in df.columns else f"'{row.parent}' as __source"
+                query = f"select {', '.join(cols)}, {source} from {row.parent}"
+                queries.append(query)
+
+            sql = f"create or replace view {self.qualified_name} as {' union all '.join(queries)}"
+            sql = fix_sql(sql)
+            Logger.debug("view", extra={"job": self, "sql": sql})
+            self.spark.sql(sql)
+
+        else:
+            assert dep_df.count() == 1, "only one dependency allowed"
+
+            parent = dep_df.collect()[0].parent
+            sql = f"select * from {parent}"
+            sql = fix_sql(sql)
+            Logger.debug("view", extra={"job": self, "sql": sql})
+
+            df = self.spark.sql(sql)
+            cdc_options = self.get_cdc_context(df)
+            self.cdc.create_or_replace_view(sql, **cdc_options)
+
+    def create_or_replace_current_view(self):
+        from py4j.protocol import Py4JJavaError
+
+        try:
+            Logger.debug("create or replace current view", extra={"job": self})
+
+            df = self.spark.sql(f"select * from {self.qualified_name}")
+
+            where_clause = "-- no where clause"
+            if "__is_current" in df.columns:
+                where_clause = "where __is_current"
+
+            sql = f"""
+            create or replace view {self.qualified_name}__current with schema evolution as
+            select
+              *
+            from
+              {self.qualified_name}
+              {where_clause}
+            """
+            # sql = fix_sql(sql)
+            # Logger.debug("current view", extra={"job": self, "sql": sql})
+            self.spark.sql(sql)
+
+        except Py4JJavaError:
+            Logger.exception("🙈", extra={"job": self})
+
+    def overwrite(self):
+        self.truncate()
+        self.run()
+
+    def overwrite_schema(self):
+        Logger.warning("overwrite schema not allowed", extra={"job": self})
+
+    def get_cdc_context(self, df: DataFrame) -> dict:
+        # if dataframe, reference is passed (BUG)
+        name = f"{self.step}_{self.topic}_{self.item}__check"
+        global_temp_view = create_or_replace_global_temp_view(name=name, df=df)
+
+        not_append = not self.mode == "append"
+        nocdc = self.change_data_capture == "nocdc"
+        order_duplicate_by = self.options.job.get_dict("order_duplicate_by") or {}
+
+        rectify = False
+        if not_append and not nocdc:
+            if not self.stream and self.mode == "update" and self.table.exists():
+                timestamp = "__valid_from" if self.change_data_capture == "scd2" else "__timestamp"
+                extra_check = f" and __timestamp > coalesce((select max({timestamp}) from {self}), cast('0001-01-01' as timestamp))"
+            else:
+                extra_check = "-- no extra check"
+
+            sql = f"""
+                select 
+                  __operation
+                from
+                  {global_temp_view}
+                where
+                  true
+                  and __operation == 'reload'
+                  {extra_check}
+                limit 
+                  1
+                """
+            sql = fix_sql(sql)
+            Logger.debug("check", extra={"job": self, "sql": sql})
+
+            check_df = self.spark.sql(sql)
+            if not check_df.isEmpty():
+                rectify = True
+                Logger.debug("rectify enabled", extra={"job": self})
+
+        context = {
+            "soft_delete": self.slowly_changing_dimension,
+            "deduplicate": self.options.job.get_boolean("deduplicate", not_append),
+            "rectify": rectify,
+            "order_duplicate_by": order_duplicate_by,
+        }
+
+        if self.slowly_changing_dimension:
+            if "__key" not in df.columns:
+                context["add_key"] = True
+
+        if self.mode == "memory":
+            context["mode"] = "complete"
+        if self.mode == "latest":
+            context["filter"] = "latest"
+
+        if self.change_data_capture == "scd2":
+            context["fix_valid_from"] = True
+
+        if nocdc:
+            if "__operation" in df.columns:
+                context["except"] = ["__operation"]
+        if nocdc and self.mode == "memory":
+            if "__operation" not in df.columns:
+                context["add_operation"] = "upsert"
+                context["except"] = ["__operation"]
+
+        if not self.stream and self.mode == "update":
+            context["filter"] = "update"
+
+        return context
+
+    def for_each_batch(self, df: DataFrame, batch: Optional[int] = None):
+        assert self.persist, f"{self.mode} not allowed"
+
+        context = self.get_cdc_context(df)
+
+        # if dataframe, reference is passed (BUG)
+        name = f"{self.step}_{self.topic}_{self.item}"
+        if batch is not None:
+            name = f"{name}__{batch}"
+        global_temp_view = create_or_replace_global_temp_view(name=name, df=df)
+        sql = f"select * from {global_temp_view}"
+
+        if self.mode == "update":
+            assert not isinstance(self.cdc, NoCDC)
+            self.cdc.update(sql, **context)
+        elif self.mode == "append":
+            assert isinstance(self.cdc, NoCDC)
+            self.cdc.append(sql, **context)
+        elif self.mode == "latest":
+            assert isinstance(self.cdc, NoCDC)
+            check_df = self.spark.sql(
+                f"""
+                select 
+                  __operation
+                from
+                  {global_temp_view}
+                where
+                  __operation <> 'reload'
+                limit 
+                  1
+                """
+            )
+            assert check_df.isEmpty(), f"{check_df.collect()[0][0]} not allowed"
+            self.cdc.complete(sql, **context)
+        else:
+            raise ValueError(f"{self.mode} - not allowed")
+
+    def create(self):
+        super().create()
+        self.create_or_replace_current_view()
+
+    def register(self):
+        super().register()
+        self.create_or_replace_current_view()
+
+    def drop(self):
+        super().drop()
+        Logger.debug("drop current view", extra={"job": self})
+        self.spark.sql(f"drop view if exists {self.qualified_name}__current")
+
+    def optimize(
+        self,
+        vacuum: Optional[bool] = True,
+        optimize: Optional[bool] = True,
+        analyze: Optional[bool] = True,
+    ):
+        if self.mode == "memory":
+            Logger.debug("memory (no optimize)", extra={"job": self})
+        else:
+            super().optimize()

fabricks/core/parsers/__init__.py

@@ -0,0 +1,12 @@
+from fabricks.core.parsers.base import PARSERS, BaseParser
+from fabricks.core.parsers.decorator import parser
+from fabricks.core.parsers.get_parser import get_parser
+from fabricks.core.parsers.types import ParserOptions
+
+__all__ = [
+    "BaseParser",
+    "get_parser",
+    "parser",
+    "ParserOptions",
+    "PARSERS",
+]

fabricks/core/parsers/base.py

@@ -0,0 +1,91 @@
+from abc import ABC
+from typing import Callable, Optional, final
+
+from pyspark.sql import DataFrame, SparkSession
+from pyspark.sql.functions import col, expr, from_json, lit
+from pyspark.sql.types import MapType, StringType
+
+from fabricks.core.parsers.types import ParserOptions
+from fabricks.core.utils import clean
+from fabricks.utils.path import Path
+from fabricks.utils.read.read import read
+
+
+class BaseParser(ABC):
+    def __init__(self, options: Optional[ParserOptions], file_format: str):
+        self.options = options or {}
+        self.file_format = file_format
+
+    def add_timestamp_from_file_path(self, df: DataFrame) -> DataFrame:
+        df = df.withColumn(
+            "__split",
+            expr("split(replace(__metadata.file_path, __metadata.file_name), '/')"),
+        )
+        df = df.withColumn("__split_size", expr("size(__split)"))
+        df = df.withColumn(
+            "__timestamp",
+            expr("left(concat_ws('', slice(__split, __split_size - 4, 4), '00'), 14)"),
+        )
+        df = df.withColumn("__timestamp", expr("to_timestamp(__timestamp, 'yyyyMMddHHmmss')"))
+        df = df.drop("__split", "__split_size")
+        return df
+
+    def parse(
+        self,
+        data_path: Path,
+        schema_path: Path,
+        spark: SparkSession,
+        stream: bool,
+    ) -> DataFrame:
+        df = read(
+            stream=stream,
+            path=data_path,
+            file_format=self.file_format,
+            schema_path=schema_path,
+            options=self.options.get("read_options"),
+            spark=spark,
+        )
+        if "__timestamp" not in df.columns:
+            df = self.add_timestamp_from_file_path(df)
+        return df
+
+    @final
+    def get_data(
+        self,
+        data_path: Path,
+        schema_path: Path,
+        spark: SparkSession,
+        stream: bool,
+    ) -> DataFrame:
+        """
+        Retrieves and processes data from the specified data path using the provided schema.
+
+        Args:
+            data_path (Path): The path to the data file.
+            schema_path (Path): The path to the schema file.
+            spark (SparkSession): The SparkSession object.
+            stream (bool): Indicates whether the data should be processed as a stream.
+
+        Returns:
+            DataFrame: The processed data as a DataFrame.
+
+        Raises:
+            AssertionError: If the "__timestamp" column is missing in the DataFrame.
+            AssertionError: If the "__metadata.file_path" column is missing in the DataFrame.
+        """
+        df = self.parse(data_path=data_path, schema_path=schema_path, spark=spark, stream=stream)
+        df = df.transform(clean)
+
+        if "__rescued_data" not in df.columns:
+            df = df.withColumn("__rescued_data", lit(None).cast(StringType()))
+        df = df.withColumn("__rescued_data", from_json(col("__rescued_data"), MapType(StringType(), StringType())))  # type: ignore
+
+        assert "__timestamp" in df.columns, "__timestamp mandatory in dataframe"
+        assert df.select("__metadata.file_path"), "file_path mandatory in struct __metadata in dataframe"
+        return df
+
+    def __str__(self):
+        return f"{type(self).__name__} ({self.file_format})"
+
+
+PARSERS: dict[str, Callable[[Optional[ParserOptions]], BaseParser]] = {}

fabricks/core/parsers/decorator.py

@@ -0,0 +1,11 @@
+from typing import Callable, Optional
+
+from fabricks.core.parsers.base import PARSERS, BaseParser
+from fabricks.core.parsers.types import ParserOptions
+
+
+def parser(name: str):
+    def decorator(parser: Callable[[Optional[ParserOptions]], BaseParser]):
+        PARSERS[name] = parser
+
+    return decorator