fabricks 3.0.11__py3-none-any.whl

fabricks/core/jobs/base/processor.py

@@ -0,0 +1,249 @@
+from abc import abstractmethod
+from functools import partial
+from typing import Optional
+
+from pyspark.sql import DataFrame
+from pyspark.sql.functions import expr
+
+from fabricks.context import IS_TYPE_WIDENING, IS_UNITY_CATALOG, SECRET_SCOPE
+from fabricks.context.log import DEFAULT_LOGGER
+from fabricks.core.jobs.base.exception import (
+    PostRunCheckException,
+    PostRunCheckWarning,
+    PostRunInvokeException,
+    PreRunCheckException,
+    PreRunCheckWarning,
+    PreRunInvokeException,
+    SchemaDriftException,
+    SkipRunCheckWarning,
+)
+from fabricks.core.jobs.base.invoker import Invoker
+from fabricks.utils.write import write_stream
+
+
+class Processor(Invoker):
+    def filter_where(self, df: DataFrame) -> DataFrame:
+        f = self.options.job.get("filter_where")
+
+        if f:
+            DEFAULT_LOGGER.debug(f"filter where {f}", extra={"label": self})
+            df = df.where(f"{f}")
+
+        return df
+
+    def encrypt(self, df: DataFrame) -> DataFrame:
+        encrypted_columns = self.options.job.get_list("encrypted_columns")
+        if encrypted_columns:
+            if not IS_UNITY_CATALOG:
+                from databricks.sdk.runtime import dbutils
+
+                key = dbutils.secrets.get(scope=SECRET_SCOPE, key="encryption-key")
+            else:
+                import os
+
+                key = os.environ["FABRICKS_ENCRYPTION_KEY"]
+
+            assert key, "key not found"
+
+            for col in encrypted_columns:
+                DEFAULT_LOGGER.debug(f"encrypt column: {col}", extra={"label": self})
+                df = df.withColumn(col, expr(f"aes_encrypt({col}, '{key}')"))
+
+        return df
+
+    def restore(self, last_version: Optional[str] = None, last_batch: Optional[str] = None):
+        """
+        Restores the processor to a specific version and batch.
+
+        Args:
+            last_version (Optional[str]): The last version to restore to. If None, no version restore will be performed.
+            last_batch (Optional[str]): The last batch to restore to. If None, no batch restore will be performed.
+        """
+        if self.persist:
+            if last_version is not None:
+                _last_version = int(last_version)
+                if self.table.get_last_version() > _last_version:
+                    self.table.restore_to_version(_last_version)
+
+            if last_batch is not None:
+                current_batch = int(last_batch) + 1
+                self.rm_commit(current_batch)
+
+                assert last_batch == self.table.get_property("fabricks.last_batch")
+                assert self.paths.commits.joinpath(last_batch).exists()
+
+    def _for_each_batch(self, df: DataFrame, batch: Optional[int] = None, **kwargs):
+        DEFAULT_LOGGER.debug("start (for each batch)", extra={"label": self})
+        if batch is not None:
+            DEFAULT_LOGGER.debug(f"batch {batch}", extra={"label": self})
+
+        df = self.base_transform(df)
+
+        diffs = self.get_schema_differences(df)
+        if diffs:
+            if self.schema_drift or kwargs.get("reload", False):
+                DEFAULT_LOGGER.warning("schema drifted", extra={"label": self, "diffs": diffs})
+                self.update_schema(df=df)
+
+            else:
+                only_type_widening_compatible = all(d.type_widening_compatible for d in diffs if d.status == "changed")
+                if only_type_widening_compatible and self.table.type_widening_enabled and IS_TYPE_WIDENING:
+                    self.update_schema(df=df, widen_types=True)
+                else:
+                    raise SchemaDriftException.from_diffs(str(self), diffs)
+
+        self.for_each_batch(df, batch, **kwargs)
+
+        if batch is not None:
+            self.table.set_property("fabricks.last_batch", batch)
+
+        self.table.create_restore_point()
+        DEFAULT_LOGGER.debug("end (for each batch)", extra={"label": self})
+
+    def for_each_run(self, **kwargs):
+        DEFAULT_LOGGER.debug("start (for each run)", extra={"label": self})
+
+        if self.virtual:
+            self.create_or_replace_view()
+
+        elif self.persist:
+            assert self.table.registered, f"{self} is not registered"
+
+            df = self.get_data(stream=self.stream, **kwargs)
+            assert df is not None, "no data"
+
+            partial(self._for_each_batch, **kwargs)
+
+            if self.stream:
+                DEFAULT_LOGGER.debug("use streaming", extra={"label": self})
+                write_stream(
+                    df,
+                    checkpoints_path=self.paths.checkpoints,
+                    func=self._for_each_batch,
+                    timeout=self.timeout,
+                )
+            else:
+                self._for_each_batch(df, **kwargs)
+
+        else:
+            raise ValueError(f"{self.mode} - not allowed")
+
+        DEFAULT_LOGGER.debug("end (for each run)", extra={"label": self})
+
+    def run(
+        self,
+        retry: Optional[bool] = True,
+        schedule: Optional[str] = None,
+        schedule_id: Optional[str] = None,
+        invoke: Optional[bool] = True,
+        reload: Optional[bool] = None,
+        vacuum: Optional[bool] = None,
+        optimize: Optional[bool] = None,
+        compute_statistics: Optional[bool] = None,
+    ):
+        """
+        Run the processor.
+
+        Args:
+            retry (bool, optional): Whether to retry the execution in case of failure. Defaults to True.
+            schedule (str, optional): The schedule to run the processor on. Defaults to None.
+            schedule_id (str, optional): The ID of the schedule. Defaults to None.
+            invoke (bool, optional): Whether to invoke pre-run and post-run methods. Defaults to True.
+        """
+        last_version = None
+        last_batch = None
+        exception = None
+
+        if self.persist:
+            last_version = self.table.get_property("fabricks.last_version")
+            if last_version is not None:
+                DEFAULT_LOGGER.debug(f"last version {last_version}", extra={"label": self})
+            else:
+                last_version = str(self.table.last_version)
+
+            last_batch = self.table.get_property("fabricks.last_batch")
+            if last_batch is not None:
+                DEFAULT_LOGGER.debug(f"last batch {last_batch}", extra={"label": self})
+
+        try:
+            DEFAULT_LOGGER.info("start (run)", extra={"label": self})
+
+            if reload:
+                DEFAULT_LOGGER.debug("force reload", extra={"label": self})
+
+            if invoke:
+                self.invoke_pre_run(schedule=schedule)
+
+            if not reload:
+                self.check_skip_run()
+
+            try:
+                self.check_pre_run()
+            except PreRunCheckWarning as e:
+                exception = e
+
+            self.for_each_run(schedule=schedule, reload=reload)
+
+            try:
+                self.check_post_run()
+            except PostRunCheckWarning as e:
+                exception = e
+
+            self.check_post_run_extra()
+
+            if invoke:
+                self.invoke_post_run(schedule=schedule)
+
+            if exception:
+                raise exception
+
+            if vacuum is None:
+                vacuum = self.options.job.get("vacuum", False)
+            if optimize is None:
+                optimize = self.options.job.get("optimize", False)
+            if compute_statistics is None:
+                compute_statistics = self.options.job.get("compute_statistics", False)
+
+            if vacuum or optimize or compute_statistics:
+                self.maintain(
+                    compute_statistics=compute_statistics,
+                    optimize=optimize,
+                    vacuum=vacuum,
+                )
+
+            DEFAULT_LOGGER.info("end (run)", extra={"label": self})
+
+        except SkipRunCheckWarning as e:
+            DEFAULT_LOGGER.warning("skip run", extra={"label": self})
+            raise e
+
+        except (PreRunCheckWarning, PostRunCheckWarning) as e:
+            DEFAULT_LOGGER.warning("fail to pass warning check", extra={"label": self})
+            raise e
+
+        except (PreRunInvokeException, PostRunInvokeException) as e:
+            DEFAULT_LOGGER.exception("fail to run invoker", extra={"label": self})
+            raise e
+
+        except (PreRunCheckException, PostRunCheckException) as e:
+            DEFAULT_LOGGER.exception("fail to pass check", extra={"label": self})
+            self.restore(last_version, last_batch)
+            raise e
+
+        except AssertionError as e:
+            DEFAULT_LOGGER.exception("fail to run", extra={"label": self})
+            self.restore(last_version, last_batch)
+            raise e
+
+        except Exception as e:
+            if not self.stream or not retry:
+                DEFAULT_LOGGER.exception("fail to run", extra={"label": self})
+                self.restore(last_version, last_batch)
+                raise e
+
+            else:
+                DEFAULT_LOGGER.warning("retry to run", extra={"label": self})
+                self.run(retry=False, schedule_id=schedule_id, schedule=schedule)
+
+    @abstractmethod
+    def overwrite(self) -> None: ...

fabricks/core/jobs/bronze.py

@@ -0,0 +1,395 @@
+from typing import Optional, Sequence, Union, cast
+
+from pyspark.sql import DataFrame
+from pyspark.sql.functions import expr, lit, md5
+from pyspark.sql.types import Row
+
+from fabricks.cdc.nocdc import NoCDC
+from fabricks.context import VARIABLES
+from fabricks.context.log import DEFAULT_LOGGER
+from fabricks.core.jobs.base._types import JobDependency, TBronze
+from fabricks.core.jobs.base.job import BaseJob
+from fabricks.core.parsers import BaseParser
+from fabricks.core.parsers.get_parser import get_parser
+from fabricks.core.parsers.utils import clean
+from fabricks.metastore.view import create_or_replace_global_temp_view
+from fabricks.utils.helpers import concat_ws
+from fabricks.utils.path import Path
+from fabricks.utils.read import read
+
+
+class Bronze(BaseJob):
+    def __init__(
+        self,
+        step: TBronze,
+        topic: Optional[str] = None,
+        item: Optional[str] = None,
+        job_id: Optional[str] = None,
+        conf: Optional[Union[dict, Row]] = None,
+    ):  # type: ignore
+        super().__init__(
+            "bronze",
+            step=step,
+            topic=topic,
+            item=item,
+            job_id=job_id,
+            conf=conf,
+        )
+
+    _parser: Optional[BaseParser] = None
+
+    @property
+    def stream(self) -> bool:
+        return self.mode not in ["register"]
+
+    @property
+    def schema_drift(self) -> bool:
+        return True
+
+    @property
+    def persist(self) -> bool:
+        return self.mode in ["append", "register"]
+
+    @property
+    def virtual(self) -> bool:
+        return False
+
+    @classmethod
+    def from_job_id(cls, step: str, job_id: str, *, conf: Optional[Union[dict, Row]] = None):
+        return cls(step=cast(TBronze, step), job_id=job_id, conf=conf)
+
+    @classmethod
+    def from_step_topic_item(cls, step: str, topic: str, item: str, *, conf: Optional[Union[dict, Row]] = None):
+        return cls(step=cast(TBronze, step), topic=topic, item=item, conf=conf)
+
+    @property
+    def data_path(self) -> Path:
+        uri = self.options.job.get("uri")
+        assert uri is not None, "no uri provided in options"
+        path = Path.from_uri(uri, regex=VARIABLES)
+        return path
+
+    def get_dependencies(self, *s) -> Sequence[JobDependency]:
+        dependencies = []
+
+        parents = self.options.job.get_list("parents")
+        if parents:
+            for p in parents:
+                dependencies.append(JobDependency.from_parts(self.job_id, p, "job"))
+
+        return dependencies
+
+    def register_external_table(self):
+        options = self.conf.parser_options  # type: ignore
+        if options:
+            file_format = options.get("file_format")
+        else:
+            file_format = "delta"
+
+        DEFAULT_LOGGER.debug(f"register external table ({self.data_path})", extra={"label": self})
+
+        try:
+            df = self.spark.sql(f"select * from {file_format}.`{self.data_path}`")
+            assert len(df.columns) > 1, "external table must have at least one column"
+        except Exception as e:
+            DEFAULT_LOGGER.exception("read external table failed", extra={"label": self})
+            raise e
+
+        self.spark.sql(
+            f"create table if not exists {self.qualified_name} using {file_format} location '{self.data_path}'"
+        )
+
+    def drop_external_table(self):
+        DEFAULT_LOGGER.warning("remove external table from metastore", extra={"label": self})
+        self.spark.sql(f"drop table if exists {self.qualified_name}")
+
+    def compute_statistics_external_table(self):
+        DEFAULT_LOGGER.debug("compute statistics (external table)", extra={"label": self})
+        self.spark.sql(f"analyze table {self.qualified_name} compute statistics")
+
+    def vacuum_external_table(self, retention_hours: Optional[int] = 168):
+        from delta import DeltaTable
+
+        DEFAULT_LOGGER.debug("vacuum (external table)", extra={"label": self})
+        try:
+            dt = DeltaTable.forPath(self.spark, self.data_path.string)
+            self.spark.sql("SET self.spark.databricks.delta.retentionDurationCheck.enabled = False")
+            dt.vacuum(retention_hours)
+        finally:
+            self.spark.sql("SET self.spark.databricks.delta.retentionDurationCheck.enabled = True")
+
+    def maintain_external_table(
+        self,
+        vacuum: Optional[bool] = True,
+        compute_statistics: Optional[bool] = True,
+    ):
+        DEFAULT_LOGGER.debug("maintain (external table)", extra={"label": self})
+        if vacuum:
+            self.vacuum_external_table()
+
+        if compute_statistics:
+            self.compute_statistics_external_table()
+
+    @property
+    def parser(self) -> BaseParser:
+        if not self._parser:
+            assert self.mode not in ["register"], f"{self.mode} not allowed"
+
+            name = self.options.job.get("parser")
+            assert name is not None, "parser not found"
+
+            options = self.conf.parser_options or None  # type: ignore
+            p = get_parser(name, options)
+
+            self._parser = p
+
+        return self._parser
+
+    def parse(self, stream: bool = False) -> DataFrame:
+        """
+        Parses the data based on the specified mode and returns a DataFrame.
+
+        Args:
+            stream (bool, optional): Indicates whether the data should be read as a stream. Defaults to False.
+
+        Returns:
+            DataFrame: The parsed data as a DataFrame.
+        """
+        if self.mode == "register":
+            if stream:
+                df = read(
+                    stream=stream,
+                    path=self.data_path,
+                    file_format="delta",
+                    # spark=self.spark, (BUG)
+                )
+            else:
+                df = self.spark.sql(f"select * from {self}")
+
+            # cleaning should done by parser
+            df = clean(df)
+
+        else:
+            df = self.parser.get_data(
+                stream=stream,
+                data_path=self.data_path,
+                schema_path=self.paths.schema,
+                spark=self.spark,
+            )
+
+        return df
+
+    def get_data(
+        self,
+        stream: bool = False,
+        transform: Optional[bool] = False,
+        schema_only: Optional[bool] = False,
+        **kwargs,
+    ) -> Optional[DataFrame]:
+        df = self.parse(stream)
+        df = self.filter_where(df)
+        df = self.encrypt(df)
+
+        if transform:
+            df = self.base_transform(df)
+
+        if schema_only:
+            df = df.where("1 == 2")
+
+        return df
+
+    def add_calculated_columns(self, df: DataFrame) -> DataFrame:
+        calculated_columns = self.options.job.get_dict("calculated_columns")
+
+        if calculated_columns:
+            for key, value in calculated_columns.items():
+                DEFAULT_LOGGER.debug(f"add calculated column ({key} -> {value})", extra={"label": self})
+                df = df.withColumn(key, expr(f"{value}"))
+
+        return df
+
+    def add_hash(self, df: DataFrame) -> DataFrame:
+        if "__hash" not in df.columns:
+            fields = [f"`{c}`" for c in df.columns if not c.startswith("__")]
+            DEFAULT_LOGGER.debug("add hash", extra={"label": self})
+
+            if "__operation" in df.columns:
+                fields += ["__operation == 'delete'"]
+
+            if "__source" in df.columns:
+                fields += ["__source"]
+
+            df = df.withColumn("__hash", md5(expr(f"{concat_ws(fields)}")))
+
+        return df
+
+    def add_key(self, df: DataFrame) -> DataFrame:
+        if "__key" not in df.columns:
+            fields = self.options.job.get_list("keys")
+            if fields:
+                DEFAULT_LOGGER.debug(f"add key ({', '.join(fields)})", extra={"label": self})
+
+                if "__source" in df.columns:
+                    fields = fields + ["__source"]
+
+                fields = [f"`{f}`" for f in fields]
+                df = df.withColumn("__key", md5(expr(f"{concat_ws(fields)}")))
+
+        return df
+
+    def add_source(self, df: DataFrame) -> DataFrame:
+        if "__source" not in df.columns:
+            source = self.options.job.get("source")
+            if source:
+                DEFAULT_LOGGER.debug(f"add source ({source})", extra={"label": self})
+                df = df.withColumn("__source", lit(source))
+
+        return df
+
+    def add_operation(self, df: DataFrame) -> DataFrame:
+        if "__operation" not in df.columns:
+            operation = self.options.job.get("operation")
+            if operation:
+                DEFAULT_LOGGER.debug(f"add operation ({operation})", extra={"label": self})
+                df = df.withColumn("__operation", lit(operation))
+
+            else:
+                df = df.withColumn("__operation", lit("upsert"))
+
+        return df
+
+    def base_transform(self, df: DataFrame) -> DataFrame:
+        df = df.transform(self.extend)
+        df = df.transform(self.add_calculated_columns)
+        df = df.transform(self.add_hash)
+        df = df.transform(self.add_operation)
+        df = df.transform(self.add_source)
+        df = df.transform(self.add_key)
+
+        if "__metadata" in df.columns:
+            if self.mode == "register":
+                #  https://github.com/delta-io/delta/issues/2014 (BUG)
+                df = df.withColumn(
+                    "__metadata",
+                    expr(
+                        f"""
+                        struct(
+                            concat_ws('/', '{self.data_path}', __timestamp, __operation) as file_path,
+                            __metadata.file_name as file_name,
+                            __metadata.file_size as file_size,            
+                            __metadata.file_modification_time as file_modification_time,
+                            cast(current_date() as timestamp) as inserted
+                        )
+                        """
+                    ),
+                )
+
+            else:
+                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,
+                            cast(current_date() as timestamp) as inserted
+                        )
+                        """
+                    ),
+                )
+
+        return df
+
+    def create_or_replace_view(self):
+        DEFAULT_LOGGER.warning("create or replace view not allowed", extra={"label": self})
+
+    def overwrite_schema(self, df: Optional[DataFrame] = None):
+        DEFAULT_LOGGER.warning("schema overwrite not allowed", extra={"label": self})
+
+    def get_cdc_context(self, df: DataFrame, reload: Optional[bool] = None) -> dict:
+        return {}
+
+    def for_each_batch(self, df: DataFrame, batch: Optional[int] = None, **kwargs):
+        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}__{batch}"
+        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
+
+        assert isinstance(self.cdc, NoCDC)
+        if self.mode == "append":
+            self.cdc.append(sql, **context)
+
+    def for_each_run(self, **kwargs):
+        if self.mode == "register":
+            DEFAULT_LOGGER.debug("register (no run)", extra={"label": self})
+        elif self.mode == "memory":
+            DEFAULT_LOGGER.debug("memory (no run)", extra={"label": self})
+        else:
+            super().for_each_run(**kwargs)
+
+    def create(self):
+        if self.mode == "register":
+            self.register_external_table()
+        elif self.mode == "memory":
+            DEFAULT_LOGGER.info("memory (no table nor view)", extra={"label": self})
+        else:
+            super().create()
+
+    def register(self):
+        if self.mode == "register":
+            self.register_external_table()
+        elif self.mode == "memory":
+            DEFAULT_LOGGER.info("memory (no table nor view)", extra={"label": self})
+        else:
+            super().register()
+
+    def truncate(self):
+        if self.mode == "register":
+            DEFAULT_LOGGER.info("register (no truncate)", extra={"label": self})
+        else:
+            super().truncate()
+
+    def restore(self, last_version: Optional[str] = None, last_batch: Optional[str] = None):
+        if self.mode == "register":
+            DEFAULT_LOGGER.info("register (no restore)", extra={"label": self})
+        else:
+            super().restore()
+
+    def drop(self):
+        if self.mode == "register":
+            self.drop_external_table()
+        super().drop()
+
+    def maintain(
+        self,
+        vacuum: Optional[bool] = True,
+        optimize: Optional[bool] = True,
+        compute_statistics: Optional[bool] = True,
+    ):
+        if self.mode == "register":
+            self.maintain_external_table(vacuum=vacuum, compute_statistics=compute_statistics)
+        else:
+            super().maintain(vacuum=vacuum, optimize=optimize, compute_statistics=compute_statistics)
+
+    def vacuum(self):
+        if self.mode == "memory":
+            DEFAULT_LOGGER.info("memory (no vacuum)", extra={"label": self})
+        elif self.mode == "register":
+            self.vacuum_external_table()
+        else:
+            super().vacuum()
+
+    def overwrite(self, schedule: Optional[str] = None):
+        self.truncate()
+        self.run(schedule=schedule)