fabricks 3.0.11__py3-none-any.whl

fabricks/core/jobs/base/generator.py

@@ -0,0 +1,447 @@
+from abc import abstractmethod
+from typing import Optional, Sequence, Union, cast
+
+from pyspark.sql import DataFrame
+from pyspark.sql.functions import lit
+
+from fabricks.cdc import NoCDC
+from fabricks.context.log import DEFAULT_LOGGER
+from fabricks.core.jobs.base._types import JobDependency
+from fabricks.core.jobs.base.configurator import Configurator
+from fabricks.metastore.table import SchemaDiff
+from fabricks.metastore.view import create_or_replace_global_temp_view
+
+
+class Generator(Configurator):
+    def update_dependencies(self):
+        DEFAULT_LOGGER.info("update dependencies", extra={"label": self})
+
+        deps = self.get_dependencies()
+        if deps:
+            df = self.spark.createDataFrame([d.model_dump() for d in deps])  # type: ignore
+            cdc = NoCDC("fabricks", self.step, "dependencies")
+            cdc.delete_missing(df, keys=["dependency_id"], update_where=f"job_id = '{self.job_id}'", uuid=True)
+
+    @abstractmethod
+    def get_dependencies(self) -> Sequence[JobDependency]: ...
+
+    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():
+            DEFAULT_LOGGER.info("delete schema folder", extra={"label": 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():
+            DEFAULT_LOGGER.info("delete checkpoints folder", extra={"label": 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.joinpath(str(id))
+        if path.exists():
+            DEFAULT_LOGGER.warning(f"delete commit {id}", extra={"label": 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
+        """
+        DEFAULT_LOGGER.warning("truncate", extra={"label": 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
+        """
+        if self.options.job.get("no_drop"):
+            raise ValueError("no_drop is set, cannot drop the job")
+
+        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:
+                DEFAULT_LOGGER.warning(f"{row.count} children found", extra={"label": 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.
+        """
+        ...
+
+    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
+
+            # first take from job options, then from step options
+            job_masks = self.options.table.get("masks", None)
+            step_masks = self.step_conf.get("table_options", {}).get("masks", None)
+            if job_masks is not None:
+                masks = job_masks
+            elif step_masks is not None:
+                masks = step_masks
+            else:
+                masks = None
+
+            maximum_compatibility = self.options.table.get_boolean("maximum_compatibility", False)
+
+            if maximum_compatibility:
+                default_properties = {
+                    "delta.minReaderVersion": "1",
+                    "delta.minWriterVersion": "7",
+                    "delta.columnMapping.mode": "none",
+                }
+            elif powerbi:
+                default_properties = {
+                    "delta.columnMapping.mode": "name",
+                    "delta.minReaderVersion": "2",
+                    "delta.minWriterVersion": "5",
+                }
+            else:
+                default_properties = {
+                    "delta.enableTypeWidening": "true",
+                    "delta.enableDeletionVectors": "true",
+                    "delta.columnMapping.mode": "name",
+                    "delta.minReaderVersion": "2",
+                    "delta.minWriterVersion": "5",
+                    "delta.feature.timestampNtz": "supported",
+                }
+
+            default_properties["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("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 is not None:
+                if liquid_clustering == "auto":
+                    liquid_clustering = True
+                    cluster_by = []
+
+                else:
+                    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:
+                        DEFAULT_LOGGER.debug("could not determine clustering column", extra={"label": self})
+                        liquid_clustering = False
+                        cluster_by = None
+
+            if liquid_clustering is None:
+                cluster_by = None
+                partition_by = self.options.table.get_list("partition_by")
+                if partition_by:
+                    partitioning = True
+
+            properties = None
+            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 properties is None:
+                properties = default_properties
+
+            primary_key = self.options.table.get_dict("primary_key")
+            foreign_keys = self.options.table.get_dict("foreign_keys")
+            comments = self.options.table.get_dict("comments")
+
+            # 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"), job=self)
+            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,
+                masks=masks,
+                primary_key=primary_key,
+                foreign_keys=foreign_keys,
+                comments=comments,
+                **cdc_options,
+            )
+
+        if not self.table.exists():
+            DEFAULT_LOGGER.debug("create table", extra={"label": self})
+
+            df = self.get_data(stream=self.stream, schema_only=True)
+            if df:
+                if self.stream:
+                    # add dummy stream to be sure that the writeStream will start
+                    spark = df.sparkSession
+
+                    dummy_df = 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)
+
+        else:
+            DEFAULT_LOGGER.debug("table exists, skip creation", extra={"label": self})
+
+    def _update_schema(
+        self,
+        df: Optional[DataFrame] = None,
+        overwrite: Optional[bool] = False,
+        widen_types: Optional[bool] = False,
+    ):
+        def _update_schema(df: DataFrame, batch: Optional[int] = None):
+            context = self.get_cdc_context(df, reload=True)
+            if overwrite:
+                self.cdc.overwrite_schema(df, **context)
+            else:
+                self.cdc.update_schema(df, widen_types=widen_types, **context)
+
+        if self.persist:
+            if df is not None:
+                _update_schema(df)
+
+            else:
+                df = self.get_data(stream=self.stream, schema_only=True)
+                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, widen_types: Optional[bool] = False):
+        self._update_schema(df=df, overwrite=False, widen_types=widen_types)
+
+    def overwrite_schema(self, df: Optional[DataFrame] = None):
+        self._update_schema(df=df, overwrite=True)
+
+    def get_differences_with_deltatable(self, df: Optional[DataFrame] = None):
+        if df is None:
+            df = self.get_data(stream=self.stream)
+            assert df is not None
+            df = self.base_transform(df)
+
+        context = self.get_cdc_context(df, reload=True)
+
+        return self.cdc.get_differences_with_deltatable(df, **context)
+
+    def get_schema_differences(self, df: Optional[DataFrame] = None) -> Optional[Sequence[SchemaDiff]]:
+        if df is None:
+            df = self.get_data(stream=self.stream)
+            assert df is not None
+            df = self.base_transform(df)
+
+        context = self.get_cdc_context(df, reload=True)
+
+        return self.cdc.get_schema_differences(df, **context)
+
+    def schema_drifted(self, df: Optional[DataFrame] = None) -> Optional[bool]:
+        d = self.get_schema_differences(df)
+        if d is None:
+            return None
+        return len(d) > 0
+
+    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, auto=False)
+                else:
+                    self.table.enable_liquid_clustering(auto=True)
+        else:
+            DEFAULT_LOGGER.debug("could not enable liquid clustering", extra={"label": self})

fabricks/core/jobs/base/invoker.py

@@ -0,0 +1,206 @@
+import json
+from typing import Optional
+
+from pyspark.sql import DataFrame
+
+from fabricks.context import PATH_RUNTIME
+from fabricks.context.log import DEFAULT_LOGGER
+from fabricks.core.jobs.base.checker import Checker
+from fabricks.core.jobs.base.exception import PostRunInvokeException, PreRunInvokeException
+from fabricks.core.jobs.get_schedule import get_schedule
+from fabricks.utils.path import Path
+
+
+class Invoker(Checker):
+    def invoke(self, schedule: Optional[str] = None, **kwargs):
+        return self._invoke_job(
+            position="run",
+            schedule=schedule,
+            **kwargs,
+        )  # kwargs and return needed for get_data in gold
+
+    def invoke_pre_run(self, schedule: Optional[str] = None):
+        self._invoke_job(position="pre_run", schedule=schedule)
+        self._invoke_step(position="pre_run", schedule=schedule)
+
+    def invoke_post_run(self, schedule: Optional[str] = None):
+        self._invoke_job(position="post_run", schedule=schedule)
+        self._invoke_step(position="post_run", schedule=schedule)
+
+    def _invoke_job(self, position: str, schedule: Optional[str] = None, **kwargs):
+        invokers = self.options.invokers.get_list(position)
+        if position == "run":
+            invokers = invokers if len(invokers) > 0 else [{}]  # run must work even without run invoker options
+
+        errors = []
+
+        if invokers:
+            for i, invoker in enumerate(invokers):
+                DEFAULT_LOGGER.debug(f"invoke ({i}, {position})", extra={"label": self})
+                try:
+                    path = kwargs.get("path")
+                    if path is None:
+                        notebook = invoker.get("notebook")
+                        assert notebook, "notebook mandatory"
+                        path = PATH_RUNTIME.joinpath(notebook)
+
+                    assert path is not None, "path mandatory"
+
+                    arguments = invoker.get("arguments") or {}
+                    timeout = invoker.get("timeout")
+
+                    schema_only = kwargs.get("schema_only")
+                    if schema_only is not None:
+                        arguments["schema_only"] = schema_only
+
+                    if len(invokers) == 1 and position == "run":
+                        return self._run_notebook(
+                            path=path,
+                            arguments=arguments,
+                            timeout=timeout,
+                            schedule=schedule,
+                        )
+                    else:
+                        self._run_notebook(
+                            path=path,
+                            arguments=arguments,
+                            timeout=timeout,
+                            schedule=schedule,
+                        )
+
+                except Exception as e:
+                    DEFAULT_LOGGER.warning(f"fail to run invoker ({i}, {position})", extra={"label": self})
+
+                    if position == "pre_run":
+                        errors.append(PreRunInvokeException(e))
+                    elif position == "post_run":
+                        errors.append(PostRunInvokeException(e))
+                    else:
+                        errors.append(e)
+
+        if errors:
+            raise Exception(errors)
+
+    def _invoke_step(self, position: str, schedule: Optional[str] = None):
+        invokers = self.step_conf.get("invoker_options", {}).get(position, [])
+
+        errors = []
+
+        if invokers:
+            for i, invoker in enumerate(invokers):
+                DEFAULT_LOGGER.debug(f"invoke by step ({i}, {position})", extra={"label": self})
+                try:
+                    notebook = invoker.get("notebook")
+                    assert notebook, "notebook mandatory"
+                    path = PATH_RUNTIME.joinpath(notebook)
+
+                    arguments = invoker.get("arguments", {})
+                    timeout = invoker.get("timeout")
+
+                    self._run_notebook(
+                        path=path,
+                        arguments=arguments,
+                        timeout=timeout,
+                        schedule=schedule,
+                    )
+
+                except Exception as e:
+                    DEFAULT_LOGGER.warning(f"fail to run invoker by step ({i}, {position})", extra={"label": self})
+
+                    if position == "pre_run":
+                        errors.append(PreRunInvokeException(e))
+                    elif position == "post_run":
+                        errors.append(PostRunInvokeException(e))
+                    else:
+                        errors.append(e)
+
+        if errors:
+            raise Exception(errors)
+
+    def _run_notebook(
+        self,
+        path: Path,
+        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.
+
+        """
+        from databricks.sdk.runtime import dbutils
+
+        for file_format in [None, ".py", ".ipynb"]:
+            path_with_file_format = path.append(file_format) if file_format else path
+            if path_with_file_format.exists():
+                path = path_with_file_format
+                break
+
+        if timeout is None:
+            timeout = self.timeout
+
+        assert timeout is not None
+
+        variables = None
+        if schedule is not None:
+            variables = get_schedule(name=schedule).get("options", {}).get("variables", {})
+
+        if variables is None:
+            variables = {}
+
+        if arguments is None:
+            arguments = {}
+
+        return dbutils.notebook.run(
+            path=path.get_notebook_path(),  # type: ignore
+            timeout_seconds=timeout,  # type: ignore
+            arguments={  # type: ignore
+                "step": self.step,
+                "topic": self.topic,
+                "item": self.item,
+                **arguments,
+                "job_options": json.dumps(self.options.job.options),
+                "schedule_variables": json.dumps(variables),
+            },
+        )
+
+    def extend_job(self, df: DataFrame) -> DataFrame:
+        from fabricks.core.extenders import get_extender
+
+        extenders = self.options.extenders
+        for e in extenders:
+            name = e.get("extender")
+            DEFAULT_LOGGER.debug(f"extend ({name})", extra={"label": self})
+            arguments = e.get("arguments") or {}
+
+            extender = get_extender(name)
+            df = extender(df, **arguments)
+
+        return df
+
+    def extend_step(self, df: DataFrame) -> DataFrame:
+        from fabricks.core.extenders import get_extender
+
+        extenders = self.step_conf.get("extender_options", {})
+        for e in extenders:
+            name = e.get("extender")
+            DEFAULT_LOGGER.debug(f"extend by step ({name})", extra={"label": self})
+            arguments = e.get("arguments", {})
+
+            extender = get_extender(name)
+            df = extender(df, **arguments)
+
+        return df
+
+    def extend(self, df: DataFrame) -> DataFrame:
+        df = self.extend_job(df)
+        df = self.extend_step(df)
+        return df

fabricks/core/jobs/base/job.py

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