fabricks 2024.7.1.5__py3-none-any.whl

fabricks/core/jobs/bronze.py

@@ -0,0 +1,333 @@
+from typing import Optional, cast
+
+from pyspark.sql import DataFrame, Row
+from pyspark.sql.functions import expr, lit, md5
+
+from fabricks.cdc.nocdc import NoCDC
+from fabricks.context import VARIABLES
+from fabricks.context.log import Logger
+from fabricks.core.jobs.base.job import BaseJob
+from fabricks.core.jobs.base.types import TBronze
+from fabricks.core.parsers import BaseParser
+from fabricks.core.parsers.get_parser import get_parser
+from fabricks.core.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
+    ):  # type: ignore
+        super().__init__(
+            "bronze",
+            step=step,
+            topic=topic,
+            item=item,
+            job_id=job_id,
+        )
+
+    _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):
+        return cls(step=cast(TBronze, step), job_id=job_id)
+
+    @classmethod
+    def from_step_topic_item(cls, step: str, topic: str, item: str):
+        return cls(step=cast(TBronze, step), topic=topic, item=item)
+
+    @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, df: Optional[DataFrame] = None) -> Optional[DataFrame]:
+        dependencies = []
+        parents = self.options.job.get_list("parents")
+        if parents:
+            for p in parents:
+                dependencies.append(Row(self.job_id, p, "job"))
+        if dependencies:
+            df = self.spark.createDataFrame(dependencies, schema=["job_id", "parent", "origin"])
+            df = df.transform(self.add_dependency_details)
+            return df
+
+    def register_external_table(self):
+        options = self.conf.parser_options  # type: ignore
+        if options:
+            file_format = options.get("file_format")
+        else:
+            file_format = "delta"
+
+        Logger.debug(f"register external table ({self.data_path})", extra={"job": self})
+        self.spark.sql(
+            f"create table if not exists {self.qualified_name} using {file_format} location '{self.data_path}'"
+        )
+
+    def drop_external_table(self):
+        Logger.debug("drop external table", extra={"job": self})
+        self.spark.sql(f"drop table if exists {self.qualified_name}")
+
+    def optimize_external_table(
+        self,
+        vacuum: Optional[bool] = True,
+        analyze: Optional[bool] = True,
+    ):
+        Logger.debug("optimize external table", extra={"job": self})
+        if vacuum:
+            from delta import DeltaTable
+
+            dt = DeltaTable.forPath(self.spark, self.data_path.string)
+            retention_days = 7
+            Logger.debug(f"{self.data_path} - vacuum table (removing files older than {retention_days} days)")
+            try:
+                self.spark.sql("SET self.spark.databricks.delta.retentionDurationCheck.enabled = False")
+                dt.vacuum(retention_days * 24)
+            finally:
+                self.spark.sql("SET self.spark.databricks.delta.retentionDurationCheck.enabled = True")
+
+        if analyze:
+            Logger.debug(f"{self.data_path} - compute delta statistics")
+            self.spark.sql(f"analyze table delta.`{self.data_path}` compute delta statistics")
+
+    @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 done in 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: bool = False) -> DataFrame:
+        df = self.parse(stream)
+        df = self.filter_where(df)
+        df = self.encrypt(df)
+        if transform:
+            df = self.base_transform(df)
+        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():
+                Logger.debug(f"add calculated column ({key} -> {value})", extra={"job": 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("__")]
+            Logger.debug("add hash", extra={"job": 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:
+                Logger.debug(f"add key ({', '.join(fields)})", extra={"job": 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:
+                Logger.debug(f"add source ({source})", extra={"job": 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:
+                Logger.debug(f"add operation ({operation})", extra={"job": 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.extender)
+        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):
+        Logger.warning("create or replace view not allowed", extra={"job": self})
+
+    def overwrite_schema(self):
+        Logger.warning("schema overwrite not allowed", extra={"job": self})
+
+    def get_cdc_context(self, df: DataFrame) -> dict:
+        return {}
+
+    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}__{batch}"
+        global_temp_view = create_or_replace_global_temp_view(name=name, df=df)
+        sql = f"select * from {global_temp_view}"
+
+        assert isinstance(self.cdc, NoCDC)
+        if self.mode == "append":
+            self.cdc.append(sql, **context)
+
+    def for_each_run(self, schedule: Optional[str] = None):
+        if self.mode == "register":
+            Logger.info("register (no run)", extra={"job": self})
+        elif self.mode == "memory":
+            Logger.info("memory (no run)", extra={"job": self})
+        else:
+            super().for_each_run(schedule=schedule)
+
+    def create(self):
+        if self.mode == "register":
+            self.register_external_table()
+        elif self.mode == "memory":
+            Logger.info("memory (no table nor view)", extra={"job": self})
+        else:
+            super().create()
+
+    def register(self):
+        if self.mode == "register":
+            self.register_external_table()
+        elif self.mode == "memory":
+            Logger.info("memory (no table nor view)", extra={"job": self})
+        else:
+            super().register()
+
+    def truncate(self):
+        if self.mode == "register":
+            Logger.info("register (no truncate)", extra={"job": self})
+        else:
+            super().truncate()
+
+    def restore(self):
+        if self.mode == "register":
+            Logger.info("register (no restore)", extra={"job": self})
+        else:
+            super().restore()
+
+    def drop(self):
+        if self.mode == "register":
+            self.drop_external_table()
+        super().drop()
+
+    def optimize(
+        self,
+        vacuum: Optional[bool] = True,
+        optimize: Optional[bool] = True,
+        analyze: Optional[bool] = True,
+    ):
+        if self.mode == "memory":
+            Logger.info("memory (no optimize)", extra={"job": self})
+        elif self.mode == "register":
+            self.optimize_external_table(vacuum, analyze)
+        else:
+            super().optimize()
+
+    def overwrite(self):
+        self.truncate()
+        self.run()

fabricks/core/jobs/get_job.py

@@ -0,0 +1,126 @@
+from typing import Optional, cast, overload
+
+from pyspark.sql import Row
+
+from fabricks.core.jobs.base.job import BaseJob
+from fabricks.core.jobs.base.types import Bronzes, Golds, Silvers, TBronze, TGold, TSilver
+from fabricks.core.jobs.get_job_id import get_job_id
+
+
+@overload
+def get_job(step: str, *, job_id: str) -> BaseJob: ...
+
+
+@overload
+def get_job(step: str, *, topic: str, item: str) -> BaseJob: ...
+
+
+@overload
+def get_job(*, row: Row) -> BaseJob: ...
+
+
+@overload
+def get_job(*, job: str) -> BaseJob: ...
+
+
+def get_job(
+    step: Optional[str] = None,
+    topic: Optional[str] = None,
+    item: Optional[str] = None,
+    job_id: Optional[str] = None,
+    job: Optional[str] = None,
+    row: Optional[Row] = None,
+) -> BaseJob:
+    """
+    Retrieve a job based on the provided parameters.
+
+    Args:
+        step (Optional[str]): The step of the job.
+        topic (Optional[str]): The topic of the job.
+        item (Optional[str]): The item of the job.
+        job_id (Optional[str]): The ID of the job.
+        job (Optional[str]): The job string.
+        row (Optional[Row]): The row object containing job information.
+
+    Returns:
+        BaseJob: The retrieved job.
+
+    Raises:
+        ValueError: If the required parameters are not provided.
+
+    """
+    if row:
+        if "step" in row and "topic" in row and "item" in row:
+            j = _get_job(step=row.step, topic=row.topic, item=row.item)
+        elif "step" in row and "job_id" in row:
+            j = get_job(step=row.step, job_id=row.job_id)
+        elif "job" in row:
+            parts = row.job.split(".")
+            s = parts[0]
+            job_id = get_job_id(job=row.job)
+            j = _get_job(step=s, job_id=job_id)
+        else:
+            raise ValueError("step, topic, item or step, job_id or job mandatory")
+
+    elif job:
+        parts = job.split(".")
+        s = parts[0]
+        job_id = get_job_id(job=job)
+        j = _get_job(step=s, job_id=job_id)
+
+    elif job_id:
+        assert step, "step mandatory"
+        j = _get_job(step=step, job_id=job_id)
+
+    else:
+        assert step, "step mandatory"
+        assert topic, "topic mandatory"
+        assert item, "item mandatory"
+        j = _get_job(step=step, topic=topic, item=item)
+
+    return j
+
+
+def _get_job(
+    step: str,
+    topic: Optional[str] = None,
+    item: Optional[str] = None,
+    job_id: Optional[str] = None,
+):
+    if step in Bronzes:
+        from fabricks.core.jobs.bronze import Bronze
+
+        step = cast(TBronze, step)
+        if job_id is not None:
+            job = Bronze.from_job_id(step=step, job_id=job_id)
+        else:
+            assert topic
+            assert item
+            job = Bronze.from_step_topic_item(step=step, topic=topic, item=item)
+
+    elif step in Silvers:
+        from fabricks.core.jobs.silver import Silver
+
+        step = cast(TSilver, step)
+        if job_id is not None:
+            job = Silver.from_job_id(step=step, job_id=job_id)
+        else:
+            assert topic
+            assert item
+            job = Silver.from_step_topic_item(step=step, topic=topic, item=item)
+
+    elif step in Golds:
+        from fabricks.core.jobs.gold import Gold
+
+        step = cast(TGold, step)
+        if job_id is not None:
+            job = Gold.from_job_id(step=step, job_id=job_id)
+        else:
+            assert topic
+            assert item
+            job = Gold.from_step_topic_item(step=step, topic=topic, item=item)
+
+    else:
+        raise ValueError(f"{step} not found")
+
+    return job

fabricks/core/jobs/get_job_conf.py

@@ -0,0 +1,115 @@
+from typing import Optional, cast, overload
+
+from databricks.sdk.runtime import spark
+from pyspark.sql import Row
+
+from fabricks.context import IS_LIVE
+from fabricks.core.jobs.base.types import Bronzes, Golds, JobConf, Silvers, TBronze, TGold, TSilver, TStep
+
+
+@overload
+def get_job_conf(step: TStep, *, job_id: str) -> JobConf: ...
+
+
+@overload
+def get_job_conf(step: TStep, *, topic: str, item: str) -> JobConf: ...
+
+
+def _get_job_conf(step: TStep, row: Row) -> JobConf:
+    options = row["options"].asDict() if row["options"] else None
+    table_options = row["table_options"].asDict() if row["table_options"] else None
+    check_options = row["check_options"].asDict() if row["check_options"] else None
+    spark_options = row["spark_options"].asDict() if row["spark_options"] else None
+    invoker_options = row["invoker_options"].asDict() if row["invoker_options"] else None
+
+    if step in Bronzes:
+        from fabricks.core.jobs.base.types import JobConfBronze
+
+        assert options is not None, "no option"
+        parser_options = row["parser_options"].asDict() if row["parser_options"] else None
+        step = cast(TBronze, step)
+        return JobConfBronze(
+            job_id=row["job_id"],
+            topic=row["topic"],
+            item=row["item"],
+            step=step,
+            options=options,
+            parser_options=parser_options,
+            table_options=table_options,
+            check_options=check_options,
+            invoker_options=invoker_options,
+            spark_options=spark_options,
+            tags=row["tags"],
+        )
+
+    elif step in Silvers:
+        from fabricks.core.jobs.base.types import JobConfSilver
+
+        assert options is not None, "no option"
+        step = cast(TSilver, step)
+        return JobConfSilver(
+            job_id=row["job_id"],
+            topic=row["topic"],
+            item=row["item"],
+            step=step,
+            options=options,
+            table_options=table_options,
+            check_options=check_options,
+            invoker_options=invoker_options,
+            spark_options=spark_options,
+            tags=row["tags"],
+        )
+
+    elif step in Golds:
+        from fabricks.core.jobs.base.types import JobConfGold
+
+        assert options is not None, "no option"
+        step = cast(TGold, step)
+        return JobConfGold(
+            job_id=row["job_id"],
+            topic=row["topic"],
+            item=row["item"],
+            step=step,
+            options=options,
+            table_options=table_options,
+            check_options=check_options,
+            invoker_options=invoker_options,
+            spark_options=spark_options,
+            tags=row["tags"],
+        )
+
+    else:
+        raise ValueError(f"{step} not found")
+
+
+def get_job_conf(
+    step: TStep,
+    job_id: Optional[str] = None,
+    topic: Optional[str] = None,
+    item: Optional[str] = None,
+) -> JobConf:
+    if IS_LIVE:
+        from fabricks.core.steps import get_step
+
+        s = get_step(step=step)
+        if topic:
+            df = s.get_jobs(topic=topic)
+        else:
+            df = s.get_jobs()
+    else:
+        df = spark.sql(f"select * from fabricks.{step}_jobs")
+
+    assert df, f"{step} not found"
+
+    if job_id:
+        try:
+            row = df.where(f"job_id == '{job_id}'").collect()[0]
+        except IndexError:
+            raise ValueError(f"job not found ({step}, {job_id})")
+    else:
+        try:
+            row = df.where(f"topic == '{topic}' and item == '{item}'").collect()[0]
+        except IndexError:
+            raise ValueError(f"job not found ({step}, {topic}, {item})")
+
+    return _get_job_conf(step=step, row=row)

fabricks/core/jobs/get_job_id.py

@@ -0,0 +1,26 @@
+from typing import Optional, overload
+
+from fabricks.utils.helpers import md5
+
+
+@overload
+def get_job_id(step: str, topic: str, item: str) -> str: ...
+
+
+@overload
+def get_job_id(*, job: str) -> str: ...
+
+
+def get_job_id(
+    step: Optional[str] = None,
+    topic: Optional[str] = None,
+    item: Optional[str] = None,
+    job: Optional[str] = None,
+) -> str:
+    if not job:
+        assert step
+        assert topic
+        assert item
+        job = f"{step}.{topic}_{item}"
+
+    return md5(job)

fabricks/core/jobs/get_jobs.py

@@ -0,0 +1,89 @@
+from dataclasses import dataclass
+from typing import List, Optional, TypedDict, Union
+
+from databricks.sdk.runtime import spark
+from pyspark.sql import DataFrame, Row
+from pyspark.sql.functions import expr
+
+from fabricks.context import IS_LIVE, PATHS_RUNTIME
+from fabricks.core.jobs.base.job import BaseJob
+from fabricks.core.jobs.base.types import Modes, TStep
+from fabricks.core.jobs.get_job import get_job
+from fabricks.utils.helpers import concat_dfs, run_in_parallel
+from fabricks.utils.path import Path
+from fabricks.utils.read import read_yaml
+from fabricks.utils.schema import get_schema_for_type
+
+
+class GenericOptions(TypedDict):
+    mode: Modes
+
+
+@dataclass
+class JobConfGeneric:
+    step: TStep
+    job_id: str
+    topic: str
+    item: str
+    options: GenericOptions
+
+
+def _get_job(row: Row):
+    return get_job(row=row)
+
+
+def _get_jobs() -> DataFrame:
+    if IS_LIVE:
+        schema = get_schema_for_type(JobConfGeneric)
+
+        def _read_yaml(path: Path):
+            df = read_yaml(path, root="job", schema=schema)
+            if df:
+                df = df.withColumn("job_id", expr("md5(concat(step,'.',topic,'_',item))"))
+                return df
+
+        dfs = run_in_parallel(_read_yaml, list(PATHS_RUNTIME.values()))
+        df = concat_dfs(dfs)
+
+    else:
+        df = spark.sql("select * from fabricks.jobs")
+
+    return df
+
+
+def get_jobs(df: Optional[DataFrame] = None, convert: Optional[bool] = False) -> Union[List[BaseJob], DataFrame]:
+    """
+    Retrieves a list of jobs or a DataFrame containing job information.
+
+    Args:
+        df (Optional[DataFrame]): Optional DataFrame containing job information.
+        convert (Optional[bool]): Flag indicating whether to convert the DataFrame to a list of jobs.
+
+    Returns:
+        Union[List[BaseJob], DataFrame]: If `convert` is False, returns a list of BaseJob objects.
+                                         If `convert` is True, returns a DataFrame with selected columns.
+
+    Raises:
+        ValueError: If the DataFrame does not contain the required columns.
+
+    """
+    if not convert:
+        return _get_jobs()
+
+    else:
+        if df is None:
+            df = _get_jobs()
+        else:
+            if "step" in df.columns and "topic" in df.columns and "item" in df.columns:
+                df = df.select("step", "topic", "item")
+            elif "step" in df.columns and "job_id" in df.columns:
+                df = df.select("step", "job_id")
+            elif "job" in df.columns:
+                df = df.select("job")
+            else:
+                raise ValueError("step, topic, item or step, job_id or job mandatory")
+
+        assert df
+
+        jobs = run_in_parallel(_get_job, df)
+        return jobs