fabricks 2024.7.1.5__py3-none-any.whl

fabricks/core/jobs/base/generator.py

@@ -0,0 +1,391 @@
+from typing import Any, Optional, Union, cast
+
+from pyspark.sql import DataFrame, Row
+from pyspark.sql.functions import expr, lit
+
+from fabricks.cdc import SCD1
+from fabricks.context.log import Logger
+from fabricks.core.jobs.base.configurator import Configurator
+from fabricks.metastore.view import create_or_replace_global_temp_view
+
+
+class Generator(Configurator):
+    def update_dependencies(self):
+        Logger.info("update dependencies", extra={"job": self})
+
+        df = self.get_dependencies()
+        if df:
+            scd1 = SCD1("fabricks", self.step, "dependencies")
+            scd1.delete_missing(df, keys=["dependency_id"], update_where=f"job_id = '{self.job_id}'", uuid=True)
+
+    def add_dependency_details(self, df: DataFrame) -> DataFrame:
+        df = df.withColumn("__parent", expr("replace(parent, '__current', '')"))
+        df = df.withColumn("parent_id", expr("md5(__parent)"))
+        df = df.withColumn("dependency_id", expr("md5(concat_ws('*', job_id, parent))"))
+        df = df.drop("__parent")
+        return df
+
+    def get_dependencies(self) -> Optional[DataFrame]:
+        import re
+
+        df = self.get_data(self.stream)
+        jvm = df._sc._jvm  # type: ignore
+        explain_plan = cast(Any, jvm.PythonSQLUtils).explainString(cast(Any, df._jdf).queryExecution(), "extended")  # type: ignore
+
+        dependencies = []
+        r = re.compile(r"(?<=SubqueryAlias spark_catalog\.)[^.]*\.[^.\n]*")
+        matches = re.findall(r, explain_plan)
+        matches = list(set(matches))
+        for m in matches:
+            dependencies.append(Row(self.job_id, m, "parser"))
+        parents = self.options.job.get_list("parents") or []
+        for p in parents:
+            dependencies.append(Row(self.job_id, p, "job"))
+
+        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)
+            assert df.where("job_id == parent_id").count() == 0, "circular dependency found"
+            return df
+
+    def rm(self):
+        """
+        Removes the schema folder and checkpoints associated with the generator.
+
+        If the schema folder exists, it will be deleted. The method also calls the `rm_checkpoints` method to remove any checkpoints associated with the generator.
+        """
+        if self.paths.schema.exists():
+            Logger.info("delete schema folder", extra={"job": self})
+            self.paths.schema.rm()
+        self.rm_checkpoints()
+
+    def rm_checkpoints(self):
+        """
+        Removes the checkpoints folder if it exists.
+
+        This method checks if the checkpoints folder exists and deletes it if it does.
+        """
+        if self.paths.checkpoints.exists():
+            Logger.info("delete checkpoints folder", extra={"job": self})
+            self.paths.checkpoints.rm()
+
+    def rm_commit(self, id: Union[str, int]):
+        """
+        Remove a commit with the given ID.
+
+        Args:
+            id (Union[str, int]): The ID of the commit to remove.
+
+        Returns:
+            None
+        """
+        path = self.paths.commits.join(str(id))
+        if path.exists():
+            Logger.warning(f"delete commit {id}", extra={"job": self})
+            path.rm()
+
+    def truncate(self):
+        """
+        Truncates the job by removing all data associated with it.
+
+        This method removes the job from the system and, if the `persist` flag is set to True,
+        it also truncates the associated table.
+
+        Returns:
+            None
+        """
+        Logger.warning("truncate", extra={"job": self})
+        self.rm()
+        if self.persist:
+            self.table.truncate()
+
+    def drop(self):
+        """
+        Drops the current job and its dependencies.
+
+        This method drops the current job and its dependencies by performing the following steps:
+        1. Queries the database to check if there are any child jobs associated with the current job.
+        2. If child jobs are found, logs a warning message and prints the list of child jobs.
+        3. Drops the current job's change data capture (cdc).
+        4. Removes the current job.
+
+        Note: This method handles any exceptions that occur during the process.
+
+        Returns:
+                None
+        """
+        try:
+            row = self.spark.sql(
+                f"""
+                select
+                    count(*) as count,
+                    array_join(sort_array(collect_set(j.job)), ', \n') as children
+                from
+                    fabricks.dependencies d
+                    inner join fabricks.jobs j on d.job_id = j.job_id
+                where
+                    parent like '{self}'
+                """
+            ).collect()[0]
+            if cast(int, row.count) > 0:
+                Logger.warning(f"{row.count} children found", extra={"job": self, "content": row.children})
+        except Exception:
+            pass
+        self.cdc.drop()
+        self.rm()
+
+    def create(self):
+        """
+        Creates a table or view based on the specified mode.
+
+        If `persist` is True, it creates a table by calling the `create_table` method.
+        If `virtual` is True, it creates or replaces a view by calling the `create_or_replace_view` method.
+        If neither `persist` nor `virtual` is True, it raises a ValueError.
+
+        Raises:
+            ValueError: If neither `persist` nor `virtual` is True.
+
+        """
+        if self.persist:
+            self.create_table()
+        elif self.virtual:
+            self.create_or_replace_view()
+        else:
+            raise ValueError(f"{self.mode} not allowed")
+
+    def register(self):
+        """
+        Register the job.
+
+        If `persist` is True, the job's table is registered.
+        If `virtual` is True, a view is created or replaced.
+        Otherwise, a ValueError is raised.
+
+        Raises:
+            ValueError: If `persist` and `virtual` are both False.
+
+        """
+        if self.persist:
+            self.table.register()
+        elif self.virtual:
+            self.create_or_replace_view()
+        else:
+            raise ValueError(f"{self.mode} not allowed")
+
+    def create_or_replace_view(self):
+        """
+        Creates or replaces a view.
+
+        This method is responsible for creating or replacing a view in the database.
+        It should be implemented by subclasses to define the specific logic for creating or replacing the view.
+
+        Raises:
+            NotImplementedError: This method is meant to be overridden by subclasses.
+        """
+        raise NotImplementedError()
+
+    def create_table(self):
+        def _create_table(df: DataFrame, batch: Optional[int] = 0):
+            df = self.base_transform(df)
+            cdc_options = self.get_cdc_context(df)
+
+            cluster_by = []
+            partition_by = []
+
+            powerbi = False
+            liquid_clustering = False
+            partitioning = False
+            identity = False
+
+            # first take from job options, then from step options
+            job_powerbi = self.options.table.get_boolean("powerbi", None)
+            step_powerbi = self.step_conf.get("table_options", {}).get("powerbi", None)
+            if job_powerbi is not None:
+                powerbi = job_powerbi
+            elif step_powerbi is not None:
+                powerbi = step_powerbi
+
+            if powerbi:
+                properties = {
+                    "delta.columnMapping.mode": "name",
+                    "delta.minReaderVersion": "2",
+                    "delta.minWriterVersion": "5",
+                    "fabricks.last_version": "0",
+                }
+            else:
+                properties = {
+                    "delta.enableDeletionVectors": "true",
+                    "delta.columnMapping.mode": "name",
+                    "delta.minReaderVersion": "2",
+                    "delta.minWriterVersion": "5",
+                    "delta.feature.timestampNtz": "supported",
+                    "fabricks.last_version": "0",
+                }
+
+            if "__identity" in df.columns:
+                identity = False
+            else:
+                identity = self.options.table.get_boolean("identity", False)
+
+            # first take from job options, then from step options
+            liquid_clustering_job = self.options.table.get_boolean("liquid_clustering", None)
+            liquid_clustering_step = self.step_conf.get("table_options", {}).get("liquid_clustering", None)
+            if liquid_clustering_job is not None:
+                liquid_clustering = liquid_clustering_job
+            elif liquid_clustering_step:
+                liquid_clustering = liquid_clustering_step
+
+            if liquid_clustering:
+                cluster_by = self.options.table.get_list("cluster_by") or []
+                if not cluster_by:
+                    if "__source" in df.columns:
+                        cluster_by.append("__source")
+                    if "__is_current" in df.columns:
+                        cluster_by.append("__is_current")
+                    if "__key" in df.columns:
+                        cluster_by.append("__key")
+                    elif "__hash" in df.columns:
+                        cluster_by.append("__hash")
+
+                if not cluster_by:
+                    Logger.warning("liquid clustering disabled (no clustering columns found)", extra={"job": self})
+                    liquid_clustering = False
+                    cluster_by = None
+
+            if not liquid_clustering:
+                cluster_by = None
+                partition_by = self.options.table.get_list("partition_by")
+                if partition_by:
+                    partitioning = True
+
+            if not powerbi:
+                # first take from job options, then from step options
+                if self.options.table.get_dict("properties"):
+                    properties = self.options.table.get_dict("properties")
+                elif self.step_conf.get("table_options", {}).get("properties", {}):
+                    properties = self.step_conf.get("table_options", {}).get("properties", {})
+
+            # if dataframe, reference is passed (BUG)
+            name = f"{self.step}_{self.topic}_{self.item}__init"
+            global_temp_view = create_or_replace_global_temp_view(name=name, df=df.where("1 ==2 "))
+            sql = f"select * from {global_temp_view}"
+
+            self.cdc.create_table(
+                sql,
+                identity=identity,
+                liquid_clustering=liquid_clustering,
+                cluster_by=cluster_by,
+                partitioning=partitioning,
+                partition_by=partition_by,
+                properties=properties,
+                **cdc_options,
+            )
+
+        if not self.table.exists():
+            df = self.get_data(self.stream)
+            if df:
+                if self.stream:
+                    # add dummy stream to be sure that the writeStream will start
+                    dummy_df = self.spark.readStream.table("fabricks.dummy")
+                    # __metadata is always present
+                    dummy_df = dummy_df.withColumn("__metadata", lit(None))
+                    dummy_df = dummy_df.select("__metadata")
+
+                    df = df.unionByName(dummy_df, allowMissingColumns=True)
+                    path = self.paths.checkpoints.append("__init")
+                    if path.exists():
+                        path.rm()
+
+                    query = (
+                        df.writeStream.foreachBatch(_create_table)
+                        .option("checkpointLocation", path.string)
+                        .trigger(once=True)
+                        .start()
+                    )
+                    query.awaitTermination()
+                    path.rm()
+                else:
+                    _create_table(df)
+
+                constraints = self.options.table.get_dict("constraints")
+                if constraints:
+                    for key, value in constraints.items():
+                        self.table.add_constraint(name=key, expr=value)
+
+                comment = self.options.table.get("comment")
+                if comment:
+                    self.table.add_comment(comment=comment)
+
+    def _update_schema(self, df: Optional[DataFrame] = None, overwrite: Optional[bool] = False):
+        def _update_schema(df: DataFrame, batch: Optional[int] = None):
+            if overwrite:
+                self.cdc.overwrite_schema(df)
+            else:
+                self.cdc.update_schema(df)
+
+        if self.persist:
+            if df is not None:
+                _update_schema(df)
+            else:
+                df = self.get_data(self.stream)
+                assert df is not None
+                df = self.base_transform(df)
+
+                if self.stream:
+                    path = self.paths.checkpoints.append("__schema")
+                    query = (
+                        df.writeStream.foreachBatch(_update_schema)
+                        .option("checkpointLocation", path.string)
+                        .trigger(once=True)
+                        .start()
+                    )
+                    query.awaitTermination()
+                    path.rm()
+                else:
+                    _update_schema(df)
+
+        elif self.virtual:
+            self.create_or_replace_view()
+        else:
+            raise ValueError(f"{self.mode} not allowed")
+
+    def update_schema(self, df: Optional[DataFrame] = None):
+        Logger.info("update schema", extra={"job": self})
+        self._update_schema(df=df, overwrite=False)
+
+    def overwrite_schema(self, df: Optional[DataFrame] = None):
+        Logger.info("overwrite schema", extra={"job": self})
+        self._update_schema(df=df, overwrite=True)
+
+    def enable_liquid_clustering(self):
+        df = self.table.dataframe
+        enable = False
+
+        # first take from job options, then from step options
+        enable_job = self.options.table.get_boolean("liquid_clustering", None)
+        enable_step = self.step_conf.get("table_options", {}).get("liquid_clustering", None)
+        if enable_job is not None:
+            enable = enable_job
+        elif enable_step:
+            enable = enable_step
+
+        if enable:
+            cluster_by = self.options.table.get_list("cluster_by") or []
+            if not cluster_by:
+                if "__source" in df.columns:
+                    cluster_by.append("__source")
+                if "__is_current" in df.columns:
+                    cluster_by.append("__is_current")
+                if "__key" in df.columns:
+                    cluster_by.append("__key")
+                elif "__hash" in df.columns:
+                    cluster_by.append("__hash")
+
+                if len(cluster_by) > 0:
+                    self.table.enable_liquid_clustering(cluster_by)
+                else:
+                    Logger.warning("liquid clustering not enabled (no clustering column found)", extra={"job": self})
+
+        else:
+            Logger.debug("liquid clustering not enabled", extra={"job": self})

fabricks/core/jobs/base/invoker.py

@@ -0,0 +1,119 @@
+import json
+from typing import Optional, overload
+
+from fabricks.context import PATH_RUNTIME
+from fabricks.context.log import Logger
+from fabricks.core.jobs.base.checker import Checker
+from fabricks.core.jobs.base.error import InvokerFailedException
+from fabricks.core.schedules import get_schedule
+from fabricks.utils.path import Path
+
+
+class Invoker(Checker):
+    def pre_run_invoke(self, schedule: Optional[str] = None):
+        self._job_position_invoke(position="pre_run", schedule=schedule)
+        self._step_position_invoke(position="pre_run", schedule=schedule)
+
+    def post_run_invoke(self, schedule: Optional[str] = None):
+        self._job_position_invoke(position="post_run", schedule=schedule)
+        self._step_position_invoke(position="post_run", schedule=schedule)
+
+    def _job_position_invoke(self, position: str, schedule: Optional[str] = None):
+        if self.options.invoker.get(position):
+            Logger.info(f"{position}-invoke", extra={"job": self})
+            try:
+                options = self.options.invoker.get_dict(position)
+                assert options
+
+                notebook = options.notebook  # type: ignore
+                assert notebook, "notebook mandatory"
+                path = PATH_RUNTIME.join(notebook)
+
+                arguments = options.arguments or {}  # type: ignore
+                timeout = arguments.get("timeout")
+                if timeout is None:
+                    if position == "pre_run":
+                        timeout = self.timeouts.pre_run
+                    elif position == "post_run":
+                        timeout = self.timeouts.post_run
+
+                self.invoke(path, arguments, timeout, schedule)
+            except Exception:
+                raise InvokerFailedException(position)
+
+    def _step_position_invoke(self, position: str, schedule: Optional[str] = None):
+        if self.step_conf.get("options", {}).get(position, None):
+            Logger.info(f"{self.step} - {position}-invoke")
+            try:
+                options = self.step_conf.get("options", {}).get(position, None)
+                assert options
+
+                notebook = options.get("notebook")  # type: ignore
+                assert notebook, "notebook mandatory"
+                path = PATH_RUNTIME.join(notebook)
+
+                arguments = options.get("arguments", {})  # type: ignore
+                timeout = arguments.get("timeout")
+                if timeout is None:
+                    if position == "pre_run":
+                        timeout = self.timeouts.pre_run
+                    elif position == "post_run":
+                        timeout = self.timeouts.post_run
+
+                self.invoke(path, arguments, timeout, schedule)
+            except Exception:
+                raise InvokerFailedException(position)
+
+    @overload
+    def invoke(self, path: Path, arguments: dict, timeout: Optional[int] = None, schedule: Optional[str] = None): ...
+
+    @overload
+    def invoke(self, *, schedule: Optional[str] = None): ...
+
+    def invoke(
+        self,
+        path: Optional[Path] = None,
+        arguments: Optional[dict] = None,
+        timeout: Optional[int] = None,
+        schedule: Optional[str] = None,
+    ):
+        """
+        Invokes a notebook job.
+
+        Args:
+            path (Optional[Path]): The path to the notebook file. If not provided, it will be retrieved from the invoker options.
+            arguments (Optional[dict]): Additional arguments to pass to the notebook job. If not provided, it will be retrieved from the invoker options.
+            schedule (Optional[str]): The schedule for the job. If provided, schedule variables will be retrieved.
+
+        Raises:
+            AssertionError: If the specified path does not exist.
+
+        """
+        if path is None:
+            notebook = self.options.invoker.get_dict("notebook")
+            path = PATH_RUNTIME.join(notebook)
+        assert path.exists(), f"{path} not found"
+
+        if arguments is None:
+            arguments = self.options.invoker.get_dict("arguments") or {}
+
+        if schedule is not None:
+            variables = get_schedule(schedule).select("options.variables").collect()[0][0]
+        else:
+            variables = {}
+
+        if timeout is None:
+            timeout = self.timeouts.job
+
+        self.dbutils.notebook.run(
+            path.get_notebook_path(),
+            timeout,
+            {
+                "step": self.step,
+                "topic": self.topic,
+                "item": self.item,
+                **arguments,
+                "job_options": json.dumps(self.options.job.options),
+                "schedule_variables": json.dumps(variables),
+            },
+        )

fabricks/core/jobs/base/job.py

@@ -0,0 +1,5 @@
+from fabricks.core.jobs.base.processor import Processor
+
+
+class BaseJob(Processor):
+    pass