fabricks 2024.7.1.5__py3-none-any.whl

fabricks/core/steps/base.py

@@ -0,0 +1,282 @@
+from typing import Optional, Union, cast
+
+from databricks.sdk.runtime import spark
+from pyspark.sql import DataFrame, Row
+from pyspark.sql.functions import expr, md5
+
+from fabricks.cdc import SCD1
+from fabricks.context import CONF_RUNTIME, PATHS_RUNTIME, PATHS_STORAGE, STEPS
+from fabricks.context.log import Logger
+from fabricks.core.jobs.base.types import Bronzes, Golds, Silvers, TStep
+from fabricks.core.jobs.get_job import get_job
+from fabricks.core.steps.get_step_conf import get_step_conf
+from fabricks.core.steps.types import Timeouts
+from fabricks.metastore.database import Database
+from fabricks.metastore.table import Table
+from fabricks.utils.helpers import concat_dfs, run_in_parallel
+from fabricks.utils.read.read_yaml import read_yaml
+from fabricks.utils.schema import get_schema_for_type
+
+
+class BaseStep:
+    def __init__(self, step: Union[TStep, str]):
+        self.name = cast(str, step)
+
+        if self.name in Bronzes:
+            self.extend = "bronze"
+        elif self.name in Silvers:
+            self.extend = "silver"
+        elif self.name in Golds:
+            self.extend = "gold"
+        else:
+            raise ValueError(self.name, "does not extend a default job")
+
+        _storage = PATHS_STORAGE.get(self.name)
+        assert _storage
+        _runtime = PATHS_RUNTIME.get(self.name)
+        assert _runtime
+
+        self.spark = spark
+        self.storage = _storage
+        self.runtime = _runtime
+        self.database = Database(self.name)
+
+    _conf: Optional[dict] = None
+    _options: Optional[dict] = None
+
+    _workers: Optional[int] = None
+    _timeouts: Optional[Timeouts] = None
+
+    @property
+    def workers(self):
+        if not self._workers:
+            w = self.options.get("workers")
+            if w is None:
+                w = CONF_RUNTIME.get("options", {}).get("workers")
+            assert w is not None
+            self._workers = cast(int, w)
+        return self._workers
+
+    def _get_timeout(self, what: str) -> int:
+        t = self.options.get("timeouts", {}).get(what, None)
+        if t is None:
+            t = CONF_RUNTIME.get("options", {}).get("timeouts", {}).get(what)
+        assert t is not None
+        return int(t)
+
+    @property
+    def timeouts(self) -> Timeouts:
+        if not self._timeouts:
+            self._timeouts = Timeouts(
+                job=self._get_timeout("job"),
+                step=self._get_timeout("step"),
+            )
+        return self._timeouts
+
+    @property
+    def conf(self) -> dict:
+        if not self._conf:
+            _conf = [s for s in STEPS if s.get("name") == self.name][0]
+            assert _conf is not None
+            self._conf = cast(dict[str, str], _conf)
+        return self._conf
+
+    @property
+    def options(self) -> dict:
+        if not self._options:
+            o = self.conf.get("options")
+            assert o is not None
+            self._options = cast(dict[str, str], o)
+        return self._options
+
+    def drop(self):
+        Logger.warning("💣 (drop)", extra={"step": self})
+
+        fs = self.database.storage
+        assert fs
+
+        tmp = fs.join("tmp")
+        if tmp.exists():
+            tmp.rm()
+        checkpoint = fs.join("checkpoints")
+        if checkpoint.exists():
+            checkpoint.rm()
+        schema = fs.join("schemas")
+        if schema.exists():
+            schema.rm()
+
+        for t in ["jobs", "tables", "dependencies", "views"]:
+            tbl = Table("fabricks", self.name, t)
+            tbl.drop()
+
+        self.database.drop()
+
+    def create(self):
+        Logger.info("🌟 (create)", extra={"step": self})
+
+        if not self.runtime.exists():
+            Logger.warning(f"{self.name} not found in runtime ({self.runtime})")
+        else:
+            self.update()
+
+    def update(self, update_dependencies: Optional[bool] = True):
+        if not self.runtime.exists():
+            Logger.warning(f"{self.name} not found in runtime ({self.runtime})")
+        else:
+            if not self.database.exists():
+                self.database.create()
+
+            self.update_jobs()
+            self.create_jobs()
+
+            if update_dependencies:
+                self.update_dependencies()
+
+            self.update_tables()
+            self.update_views()
+
+    def get_dependencies(self) -> Optional[DataFrame]:
+        errors = []
+
+        def _get_dependencies(row: Row):
+            job = get_job(step=self.name, job_id=row["job_id"])
+            try:
+                df = job.get_dependencies()
+            except:  # noqa E722
+                errors.append(job)
+            return df
+
+        job_df = self.get_jobs()
+        if job_df:
+            dfs = run_in_parallel(_get_dependencies, job_df, workers=32)
+            for e in errors:
+                Logger.error("failed to get dependencies", extra={"step": e})
+
+            if dfs:
+                df = concat_dfs(dfs)
+                return df if not df.isEmpty() else None
+
+    def get_jobs(self, topic: Optional[str] = None) -> Optional[DataFrame]:
+        try:
+            conf = get_step_conf(self.name)
+            schema = get_schema_for_type(conf)
+
+            df = None
+            if topic:
+                df = read_yaml(self.runtime, root="job", schema=schema, file_name=topic)  # type: ignore
+
+            if not df:
+                df = read_yaml(self.runtime, root="job", schema=schema)  # type: ignore
+            elif df.isEmpty():
+                df = read_yaml(self.runtime, root="job", schema=schema)  # type: ignore
+
+            if df:
+                df = df.withColumn("job_id", md5(expr("concat(step, '.' ,topic, '_', item)")))
+
+                duplicated_df = df.groupBy("job_id", "step", "topic", "item").count().where("count > 1")
+                duplicates = ",".join(f"{row.step}.{row.topic}_{row.item}" for row in duplicated_df.collect())
+                assert duplicated_df.isEmpty(), f"duplicated job(s) ({duplicates})"
+
+                return df if not df.isEmpty() else None
+
+        except AssertionError as e:
+            Logger.exception("🙈", extra={"step": self})
+            raise e
+
+    def create_jobs(self, retry: Optional[bool] = True):
+        Logger.info("create jobs", extra={"step": self})
+
+        errors = []
+
+        def _create_job(row: Row):
+            job = get_job(step=self.name, job_id=row["job_id"])
+            try:
+                job.create()
+            except:  # noqa E722
+                errors.append(job)
+
+        df = self.get_jobs()
+        table_df = self.database.get_tables()
+        view_df = self.database.get_views()
+
+        if df:
+            if table_df:
+                table_df = table_df.withColumn("job_id", expr("md5(table)"))
+                df = df.join(table_df, "job_id", how="left_anti")
+            if view_df:
+                view_df = view_df.withColumn("job_id", expr("md5(view)"))
+                df = df.join(view_df, "job_id", how="left_anti")
+
+            run_in_parallel(_create_job, df)
+            if errors:
+                for e in errors:
+                    Logger.error("not created", extra={"job": e})
+
+                if retry:
+                    Logger.warning("retry create jobs", extra={"step": self})
+                    self.update_tables()
+                    self.update_views()
+                    self.create_jobs(retry=False)
+                else:
+                    Logger.warning("retry failed", extra={"step": self})
+        else:
+            Logger.debug("no new job", extra={"step": self})
+
+    def update_jobs(self, drop: Optional[bool] = False):
+        df = self.get_jobs()
+        if df:
+            Logger.info("update jobs", extra={"step": self})
+            if drop:
+                SCD1("fabricks", self.name, "jobs").table.drop()
+            SCD1("fabricks", self.name, "jobs").delete_missing(df, keys=["job_id"])
+        else:
+            Logger.debug("no job", extra={"step": self})
+
+    def update_tables(self):
+        df = self.database.get_tables()
+        if df:
+            Logger.debug("update tables", extra={"step": self})
+            df = df.withColumn("job_id", expr("md5(table)"))
+            SCD1("fabricks", self.name, "tables").delete_missing(df, keys=["job_id"])
+        else:
+            Logger.debug("no table", extra={"step": self})
+
+    def update_views(self):
+        df = self.database.get_views()
+        if df:
+            Logger.debug("update views", extra={"step": self})
+            df = df.withColumn("job_id", expr("md5(view)"))
+            SCD1("fabricks", self.name, "views").delete_missing(df, keys=["job_id"])
+        else:
+            Logger.debug("no view", extra={"step": self})
+
+    def update_dependencies(self):
+        df = self.get_dependencies()
+        if df:
+            Logger.debug("update dependencies", extra={"step": self})
+            df.cache()
+            SCD1("fabricks", self.name, "dependencies").delete_missing(df, keys=["dependency_id"])
+        else:
+            Logger.debug("no dependency", extra={"step": self})
+
+    def register(self, update: Optional[bool] = False, drop: Optional[bool] = False):
+        def _register(row: Row):
+            job = get_job(step=self.name, topic=row["topic"], item=row["item"])
+            job.register()
+
+        if drop:
+            spark.sql(f"drop database if exists {self.name} cascade ")
+            spark.sql(f"create database {self.name}")
+        if update:
+            self.update_jobs()
+
+        df = self.get_jobs()
+        if df:
+            table_df = self.database.get_tables()
+            if table_df:
+                df = df.join(table_df, "job_id", how="left_anti")
+        if df:
+            run_in_parallel(_register, df, workers=16)
+
+    def __str__(self):
+        return self.name

fabricks/core/steps/get_step.py

@@ -0,0 +1,10 @@
+from typing import Union
+
+from fabricks.core.jobs.base.types import Steps, TStep
+from fabricks.core.steps.base import BaseStep
+
+
+def get_step(step: Union[TStep, str]) -> BaseStep:
+    assert step in Steps, f"{step} not found"
+    _step = BaseStep(step=step)
+    return _step

fabricks/core/steps/get_step_conf.py

@@ -0,0 +1,33 @@
+from typing import Union, cast
+
+from fabricks.core.jobs.base.types import (
+    Bronzes,
+    Golds,
+    JobConfBronze,
+    JobConfGold,
+    JobConfSilver,
+    Silvers,
+    TStep,
+)
+
+
+def get_step_conf(step: Union[TStep, str]):
+    if isinstance(step, str):
+        step = cast(TStep, step)
+
+    if step in Bronzes:
+        extend = "bronze"
+    elif step in Silvers:
+        extend = "silver"
+    elif step in Golds:
+        extend = "gold"
+    else:
+        raise ValueError(f"{step} - not found")
+
+    job_conf = {
+        "bronze": JobConfBronze,
+        "silver": JobConfSilver,
+        "gold": JobConfGold,
+    }.get(extend, None)
+    assert job_conf
+    return job_conf

fabricks/core/steps/types.py

@@ -0,0 +1,7 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class Timeouts:
+    job: int
+    step: int

fabricks/core/udfs.py

@@ -0,0 +1,106 @@
+import importlib.util
+import os
+import re
+from typing import Callable, List, Optional
+
+from databricks.sdk.runtime import spark as _spark
+from pyspark.sql import SparkSession
+
+from fabricks.context import PATH_UDFS
+from fabricks.context.log import Logger
+from fabricks.core.site_packages import add_site_packages_to_path
+
+UDFS: dict[str, Callable] = {}
+
+
+def register_all_udfs():
+    """
+    Register all user-defined functions (UDFs).
+
+    This function iterates over all UDFs returned by the `get_udfs` function,
+    splits the UDF name into the function name and extension, and attempts to
+    register the UDF using the `register_udf` function. If an exception occurs
+    during registration, an error message is logged.
+
+    Returns:
+        None
+    """
+    for udf in get_udfs():
+        split = udf.split(".")
+        try:
+            register_udf(udf=split[0], extension=split[1])
+        except Exception:
+            Logger.exception(f"udf {udf} not registered")
+
+
+def get_udfs() -> List[str]:
+    files = [os.path.basename(f) for f in PATH_UDFS.walk()]
+    udfs = [f for f in files if not str(f).endswith("__init__.py") and not str(f).endswith(".requirements.txt")]
+    return udfs
+
+
+def get_extension(udf: str) -> str:
+    for u in get_udfs():
+        r = re.compile(rf"{udf}(\.py|\.sql)")
+        if re.match(r, u):
+            return u.split(".")[1]
+    raise ValueError(f"{udf} not found")
+
+
+def is_registered(udf: str, spark: Optional[SparkSession] = None) -> bool:
+    if spark is None:
+        spark = _spark
+    assert spark is not None
+
+    df = spark.sql("show functions in default")
+    df = df.where(f"function == 'spark_catalog.default.udf_{udf}'")
+    return not df.isEmpty()
+
+
+def register_udf(udf: str, extension: Optional[str] = None, spark: Optional[SparkSession] = None):
+    """
+    Register a user-defined function (UDF) in Spark.
+
+    Args:
+        udf (str): The name of the UDF to register.
+        extension (Optional[str]): The file extension of the UDF implementation file. If not provided, it will be inferred from the UDF name.
+        spark (Optional[SparkSession]): The SparkSession object. If not provided, a new SparkSession will be created.
+
+    Raises:
+        ValueError: If the UDF implementation file is not found or if the UDF name is not found.
+
+    """
+    if spark is None:
+        spark = _spark
+    assert spark is not None
+
+    if not is_registered(udf, spark):
+        if extension is None:
+            extension = get_extension(udf)
+
+        assert extension
+        path = PATH_UDFS.join(f"{udf}.{extension}")
+        if extension == "sql":
+            spark.sql(path.get_sql())
+
+        elif extension == "py":
+            assert path.exists(), f"udf not found ({path.string})"
+            spec = importlib.util.spec_from_file_location(udf, path.string)
+            assert spec, f"no valid udf found ({path.string})"
+            spec.loader.load_module()  # type: ignore
+
+            u = UDFS[udf]
+            u(spark)
+
+        else:
+            raise ValueError(f"{udf} not found")
+
+
+def udf(name: str):
+    add_site_packages_to_path()
+
+    def decorator(fn: Callable):
+        UDFS[name] = fn
+        return fn
+
+    return decorator

fabricks/core/utils.py

@@ -0,0 +1,69 @@
+from pyspark.sql import DataFrame
+from pyspark.sql.functions import length, lower
+from pyspark.sql.functions import trim as _trim
+from pyspark.sql.functions import when
+from pyspark.sql.types import DoubleType, FloatType, IntegerType
+
+
+def value_to_none(df: DataFrame) -> DataFrame:
+    cols = [name for name, dtype in df.dtypes if not name.startswith("__")]
+    for c in cols:
+        df = df.withColumn(
+            c,
+            when(length(df[f"`{c}`"].cast("string")) == 0, None)
+            .when(lower(df[f"`{c}`"].cast("string")) == "none", None)
+            .when(lower(df[f"`{c}`"].cast("string")) == "null", None)
+            .when(lower(df[f"`{c}`"].cast("string")) == "blank", None)
+            .when(lower(df[f"`{c}`"].cast("string")) == "(none)", None)
+            .when(lower(df[f"`{c}`"].cast("string")) == "(null)", None)
+            .when(lower(df[f"`{c}`"].cast("string")) == "(blank)", None)
+            .otherwise(df[f"`{c}`"]),
+        )
+    return df
+
+
+def decimal_to_float(df: DataFrame) -> DataFrame:
+    cols = [name for name, dtype in df.dtypes if dtype.startswith("decimal") and not name.startswith("__")]
+    for c in cols:
+        df = df.withColumn(c, df[f"`{c}`"].cast(FloatType()))
+    return df
+
+
+def decimal_to_double(df: DataFrame) -> DataFrame:
+    cols = [name for name, dtype in df.dtypes if dtype.startswith("decimal") and not name.startswith("__")]
+    for c in cols:
+        df = df.withColumn(c, df[f"`{c}`"].cast(DoubleType()))
+    return df
+
+
+def tinyint_to_int(df: DataFrame) -> DataFrame:
+    cols = [name for name, dtype in df.dtypes if dtype.startswith("tinyint") and not name.startswith("__")]
+    for c in cols:
+        df = df.withColumn(c, df[f"`{c}`"].cast(IntegerType()))
+    return df
+
+
+def trim(df: DataFrame) -> DataFrame:
+    cols = [name for name, dtype in df.dtypes if dtype.startswith("string") and not name.startswith("__")]
+    for c in cols:
+        df = df.withColumn(c, _trim(df[f"`{c}`"]))
+    return df
+
+
+def clean(df: DataFrame) -> DataFrame:
+    """
+    Cleans the given DataFrame by performing the following operations:
+    1. Trims whitespace from all string columns.
+    2. Converts empty strings to None.
+    3. Converts decimal values to double.
+
+    Args:
+        df (pandas.DataFrame): The DataFrame to be cleaned.
+
+    Returns:
+        pandas.DataFrame: The cleaned DataFrame.
+    """
+    df = trim(df)
+    df = value_to_none(df)
+    df = decimal_to_double(df)
+    return df

fabricks/core/views.py

@@ -0,0 +1,36 @@
+from databricks.sdk.runtime import spark
+
+from fabricks.context import PATH_VIEWS
+from fabricks.context.log import Logger
+from fabricks.utils.path import Path
+from fabricks.utils.sqlglot import fix as fix_sql
+
+
+def _create_or_replace_view(path: Path):
+    sql = path.get_sql()
+    file_name = path.get_file_name().split(".")[0]
+    sql = f"""
+    create or replace view fabricks.{file_name}
+    as
+    {sql}
+    """
+    sql = fix_sql(sql)
+    Logger.debug(f"schedule - %sql\n---\n{sql}\n---")
+
+    spark.sql(sql)
+
+
+def create_or_replace_view(name: str):
+    p = PATH_VIEWS.join(f"{name}.sql")
+    try:
+        _create_or_replace_view(p)
+    except Exception:
+        Logger.warning(f"schedule - {name} not created nor replace")
+
+
+def create_or_replace_views():
+    for p in PATH_VIEWS.walk(file_format="sql", convert=True):
+        try:
+            _create_or_replace_view(p)
+        except Exception:
+            Logger.warning(f"schedule - {p.get_file_name()} not created nor replace")

fabricks/metastore/README.md

@@ -0,0 +1,3 @@
+# BMS DNA Fabricks Metastore
+
+Metastore - Fabricks

fabricks/metastore/__init__.py

@@ -0,0 +1,5 @@
+from fabricks.metastore.database import Database
+from fabricks.metastore.table import Table
+from fabricks.metastore.view import View
+
+__all__ = ["Database", "Table", "View"]

fabricks/metastore/database.py

@@ -0,0 +1,71 @@
+from typing import Optional
+
+from databricks.sdk.runtime import spark as _spark
+from pyspark.errors.exceptions.base import AnalysisException
+from pyspark.sql import DataFrame, SparkSession
+from typing_extensions import deprecated
+
+from fabricks.context import PATHS_STORAGE
+from fabricks.context.log import Logger
+from fabricks.metastore.utils import get_tables, get_views
+from fabricks.utils.path import Path
+
+
+class Database:
+    def __init__(self, name: str, spark: Optional[SparkSession] = None):
+        self.name = name
+        storage = PATHS_STORAGE.get(self.name)
+        assert storage is not None
+        self.storage = storage
+        if spark is None:
+            spark = _spark
+        assert spark is not None
+        self.spark = spark
+
+    @property
+    @deprecated("use delta_path instead")
+    def deltapath(self) -> Path:
+        return self.storage.join("delta")
+
+    @property
+    def delta_path(self) -> Path:
+        return self.storage.join("delta")
+
+    def create(self):
+        Logger.info("🌟 (create database)", extra={"step": self})
+        self.spark.sql(f"create database if not exists {self.name};")
+
+    def drop(self, rm: Optional[bool] = True):
+        if self.exists():
+            Logger.warning("💣 (drop database)", extra={"step": self})
+            self.spark.sql(f"drop database if exists {self.name} cascade;")
+
+        if rm:
+            if self.deltapath.exists():
+                Logger.debug("🧹 (remove delta files)", extra={"step": self})
+                self.deltapath.rm()
+
+    def exists(self) -> bool:
+        try:
+            self.spark.sql(f"show tables in {self.name}")
+        # database not found
+        except AnalysisException:
+            return False
+        return True
+
+    def __str__(self):
+        return self.name
+
+    def get_tables(self) -> Optional[DataFrame]:
+        try:
+            df = get_tables(self.name)
+            return df if not df.isEmpty() else None
+        except AnalysisException:
+            return None
+
+    def get_views(self) -> Optional[DataFrame]:
+        try:
+            df = get_views(self.name)
+            return df if not df.isEmpty() else None
+        except AnalysisException:
+            return None

fabricks/metastore/pyproject.toml

@@ -0,0 +1,20 @@
+[build-system]
+requires = [ "poetry_core>=1.0.0",]
+build-backend = "poetry.core.masonry.api"
+
+[tool.poetry]
+name = "fabricks-metastore"
+version = "2024.7.1.5"
+description = "Fabricks - Metastore"
+license = "MIT"
+authors = [ "BMS DWH Team <bi_support@bmsuisse.ch>",]
+readme = "README.md"
+packages = [{include="fabricks"}]
+
+[tool.black]
+line-length = 119
+
+[tool.poetry.dependencies]
+python = ">=3.9,<4"
+"fabricks.utils" = { path = "../utils", develop = true }
+"fabricks.context" = { path = "../context", develop = true }