fabricks 3.0.11__py3-none-any.whl

fabricks/core/jobs/gold.py

@@ -0,0 +1,415 @@
+import re
+from collections.abc import Sequence
+from typing import List, Optional, Union, cast
+
+from pyspark.sql import DataFrame
+from pyspark.sql.types import Row
+from typing_extensions import deprecated
+
+from fabricks.cdc.nocdc import NoCDC
+from fabricks.context.log import DEFAULT_LOGGER
+from fabricks.core.jobs.base._types import JobDependency, TGold
+from fabricks.core.jobs.base.job import BaseJob
+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
+from fabricks.utils.sqlglot import fix, get_tables
+
+
+class Gold(BaseJob):
+    def __init__(
+        self,
+        step: TGold,
+        topic: Optional[str] = None,
+        item: Optional[str] = None,
+        job_id: Optional[str] = None,
+        conf: Optional[Union[dict, Row]] = None,
+    ):  # type: ignore
+        super().__init__(
+            "gold",
+            step=step,
+            topic=topic,
+            item=item,
+            job_id=job_id,
+            conf=conf,
+        )
+
+    _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, *, conf: Optional[Union[dict, Row]] = None):
+        return cls(step=cast(TGold, step), job_id=job_id)
+
+    @classmethod
+    def from_step_topic_item(cls, step: str, topic: str, item: str, *, conf: Optional[Union[dict, Row]] = None):
+        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"]
+
+    @property
+    def sql(self) -> str:
+        sql = self.paths.runtime.get_sql()
+        return fix(sql, keep_comments=False)
+
+    @deprecated("use sql instead")
+    def get_sql(self) -> str:
+        return self.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.sql:
+                r = re.compile(r"(?<=udf_)\w*(?=\()")
+                matches = re.findall(r, self.sql)
+                matches = set(matches)
+                matches = list(matches)
+            return matches
+
+    def register_udfs(self):
+        for u in self.get_udfs():
+            if not is_registered(u):
+                DEFAULT_LOGGER.debug(f"register udf ({u})", extra={"label": self})
+                register_udf(udf=u, spark=self.spark)
+
+    def base_transform(self, df: DataFrame) -> DataFrame:
+        df = df.transform(self.extend)
+        return df
+
+    def get_data(
+        self,
+        stream: bool = False,
+        transform: Optional[bool] = False,
+        schema_only: Optional[bool] = False,
+        **kwargs,
+    ) -> 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"):
+            invokers = self.options.invokers.get_list("run")
+            assert len(invokers) <= 1, "at most one invoker allowed when notebook is true"
+
+            global_temp_view = self.invoke(path=self.paths.runtime, schema_only=schema_only, **kwargs)
+            assert global_temp_view is not None, "global_temp_view not found"
+
+            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.sql, "sql not found"
+            self.register_udfs()
+            df = self.spark.sql(self.sql)
+
+        if transform:
+            df = self.base_transform(df)
+
+        if schema_only:
+            df = df.where("1 == 2")
+
+        return df
+
+    def create_or_replace_view(self):
+        assert self.mode == "memory", f"{self.mode} not allowed"
+
+        df = self.spark.sql(self.sql)
+        cdc_options = self.get_cdc_context(df)
+        self.cdc.create_or_replace_view(self.sql, **cdc_options)
+
+    def get_dependencies(self) -> Sequence[JobDependency]:
+        data = []
+        parents = self.options.job.get_list("parents") or []
+
+        if self.mode == "invoke":
+            dependencies = []
+        elif self.options.job.get("notebook"):
+            dependencies = self._get_notebook_dependencies()
+        else:
+            dependencies = self._get_sql_dependencies()
+
+        dependencies = [d for d in dependencies if d not in parents]
+        dependencies = [d.replace("__current", "") for d in dependencies]
+        dependencies = list(set(dependencies))
+
+        for d in dependencies:
+            data.append(JobDependency.from_parts(self.job_id, d, "parser"))
+
+        for p in parents:
+            data.append(JobDependency.from_parts(self.job_id, p, "job"))
+        return data
+
+    def _get_sql_dependencies(self) -> List[str]:
+        from fabricks.core.jobs.base._types import Steps
+
+        steps = [str(s) for s in Steps]
+        return get_tables(self.sql, allowed_databases=steps)
+
+    def _get_notebook_dependencies(self) -> List[str]:
+        import re
+
+        from fabricks.context import CATALOG
+
+        dependencies = []
+        df = self.get_data(stream=self.stream)
+
+        if df is not None:
+            explain_plan = self.spark.sql("explain extended select * from {df}", df=df).collect()[0][0]
+
+            if CATALOG is None:
+                r = re.compile(r"(?<=SubqueryAlias spark_catalog\.)[^.]*\.[^.\n]*")
+            else:
+                r = re.compile(rf"(?:(?<=SubqueryAlias spark_catalog\.)|(?<=SubqueryAlias {CATALOG}\.))[^.]*\.[^.\n]*")
+
+            matches = re.findall(r, explain_plan)
+            dependencies = list(set(matches))
+
+        return dependencies
+
+    def get_cdc_context(self, df: DataFrame, reload: Optional[bool] = None) -> dict:
+        # assume no duplicate in gold (to improve performance)
+        deduplicate = self.options.job.get_boolean("deduplicate", None)
+        # assume no reload in gold (to improve performance)
+        rectify = self.options.job.get_boolean("rectify_as_upserts", None)
+
+        add_metadata = self.options.job.get_boolean("metadata", None)
+        if add_metadata is None:
+            add_metadata = self.step_conf.get("options", {}).get("metadata", False)
+
+        context = {
+            "add_metadata": add_metadata,
+            "soft_delete": True if self.slowly_changing_dimension else None,
+            "deduplicate_key": None,
+            "deduplicate_hash": True if self.slowly_changing_dimension else None,
+            "deduplicate": False,
+            "rectify": False,
+        }
+
+        # force deduplicate
+        if deduplicate is not None:
+            context["deduplicate"] = deduplicate
+            context["deduplicate_key"] = deduplicate
+            context["deduplicate_hash"] = deduplicate
+
+        # force rectify
+        if rectify is not None:
+            context["rectify"] = rectify
+
+        # add key and hash when needed
+        if self.mode == "update" and self.change_data_capture == "nocdc":
+            if "__key" not in df.columns:
+                context["add_key"] = True
+            if "__hash" not in df.columns:
+                context["add_hash"] = True
+
+        # add key and hash when needed
+        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 self.slowly_changing_dimension:
+            if "__operation" not in df.columns:
+                # assume no duplicate hash
+                if deduplicate is None:
+                    context["deduplicate_hash"] = None
+
+                if self.mode == "update":
+                    context["add_operation"] = "reload"
+                    if rectify is None:
+                        context["rectify"] = True
+
+                else:
+                    context["add_operation"] = "upsert"
+
+        # filter to get latest data
+        if not reload:
+            if self.mode == "update" and self.change_data_capture == "scd2":
+                context["slice"] = "update"
+
+            if self.mode == "update" and self.change_data_capture == "nocdc" and "__timestamp" in df.columns:
+                context["slice"] = "update"
+
+            if self.mode == "append" and "__timestamp" in df.columns:
+                context["slice"] = "update"
+
+        if self.mode == "memory":
+            context["mode"] = "complete"
+
+        # correct __valid_from
+        if self.change_data_capture == "scd2":
+            context["correct_valid_from"] = self.options.job.get_boolean("correct_valid_from", True)
+
+        # add __timestamp
+        if self.options.job.get_boolean("persist_last_timestamp"):
+            if self.change_data_capture == "scd1":
+                if "__timestamp" not in df.columns:
+                    context["add_timestamp"] = True
+            if self.change_data_capture == "scd2":
+                if "__valid_from" not in df.columns:
+                    context["add_timestamp"] = True
+
+        if "__order_duplicate_by_asc" in df.columns:
+            context["order_duplicate_by"] = {"__order_duplicate_by_asc": "asc"}
+        elif "__order_duplicate_by_desc" in df.columns:
+            context["order_duplicate_by"] = {"__order_duplicate_by_desc": "desc"}
+
+        return context
+
+    def for_each_batch(self, df: DataFrame, batch: Optional[int] = None, **kwargs):
+        assert self.persist, f"{self.mode} not allowed"
+
+        reload = kwargs.get("reload")
+        context = self.get_cdc_context(df=df, reload=reload)
+
+        # 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, job=self)
+        sql = f"select * from {global_temp_view}"
+
+        check_df = self.spark.sql(sql)
+        if check_df.isEmpty():
+            DEFAULT_LOGGER.warning("no data", extra={"label": self})
+            return
+
+        if reload:
+            DEFAULT_LOGGER.warning("force reload", extra={"label": self})
+            self.cdc.complete(sql, **context)
+
+        elif self.mode == "update":
+            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()
+        self.check_duplicate_identity()
+
+    def for_each_run(self, **kwargs):
+        last_version = None
+        if self.options.job.get_boolean("persist_last_timestamp"):
+            last_version = self.table.get_last_version()
+
+        if self.mode == "invoke":
+            schedule = kwargs.get("schedule", None)
+            self.invoke(schedule=schedule)
+        else:
+            super().for_each_run(**kwargs)
+
+        if self.options.job.get_boolean("persist_last_timestamp"):
+            self._update_last_timestamp(last_version=last_version)
+
+    def create(self):
+        if self.mode == "invoke":
+            DEFAULT_LOGGER.info("invoke (no table nor view)", extra={"label": self})
+        else:
+            self.register_udfs()
+            super().create()
+            if self.options.job.get_boolean("persist_last_timestamp"):
+                self._update_last_timestamp(create=True)
+
+    def register(self):
+        if self.options.job.get_boolean("persist_last_timestamp"):
+            self.cdc_last_timestamp.table.register()
+
+        if self.mode == "invoke":
+            DEFAULT_LOGGER.info("invoke (no table nor view)", extra={"label": self})
+        else:
+            super().register()
+
+    def drop(self):
+        if self.options.job.get_boolean("persist_last_timestamp"):
+            self.cdc_last_timestamp.drop()
+
+        super().drop()
+
+    @property
+    def cdc_last_timestamp(self) -> NoCDC:
+        assert self.mode == "update", "persist_last_timestamp only allowed in update"
+        assert self.change_data_capture in ["scd1", "scd2"], "persist_last_timestamp only allowed in scd1 or scd2"
+
+        cdc = NoCDC(self.step, self.topic, f"{self.item}__last_timestamp")
+        return cdc
+
+    def _update_last_timestamp(self, last_version: Optional[int] = None, create: bool = False):
+        df = self.spark.sql(f"select * from {self} limit 1")
+
+        fields = []
+        if self.change_data_capture == "scd1":
+            fields.append("max(__timestamp) :: timestamp as __timestamp")
+        elif self.change_data_capture == "scd2":
+            fields.append("max(__valid_from) :: timestamp as __timestamp")
+        if "__source" in df.columns:
+            fields.append("__source")
+
+        asof = None
+        if last_version is not None:
+            asof = f"version as of {last_version}"
+
+        sql = f"select {', '.join(fields)} from {self} {asof} group by all"
+        df = self.spark.sql(sql)
+
+        if create:
+            self.cdc_last_timestamp.table.create(df)
+        else:
+            self.cdc_last_timestamp.overwrite(df)
+
+    def overwrite(self, schedule: Optional[str] = None):
+        if self.mode == "invoke":
+            DEFAULT_LOGGER.debug("invoke (no overwrite)", extra={"label": self})
+            return
+
+        elif self.mode == "memory":
+            DEFAULT_LOGGER.debug("memory (no overwrite)", extra={"label": self})
+            self.create_or_replace_view()
+            return
+
+        self.overwrite_schema()
+        self.run(reload=True, schedule=schedule)