fabricks 3.0.11__py3-none-any.whl

fabricks/core/udfs.py

@@ -0,0 +1,106 @@
+import importlib.util
+import os
+import re
+from typing import Callable, List, Optional
+
+from pyspark.sql import SparkSession
+
+from fabricks.context import CATALOG, IS_UNITY_CATALOG, PATH_UDFS, SPARK
+from fabricks.context.log import DEFAULT_LOGGER
+
+UDFS: dict[str, Callable] = {}
+
+
+def register_all_udfs(extension: Optional[str] = None):
+    """
+    Register all user-defined functions (UDFs).
+    """
+    DEFAULT_LOGGER.info("register udfs")
+
+    for udf in get_udfs(extension=extension):
+        split = udf.split(".")
+        try:
+            register_udf(udf=split[0], extension=split[1])
+        except Exception as e:
+            DEFAULT_LOGGER.exception(f"could not register udf {udf}", exc_info=e)
+
+
+def get_udfs(extension: Optional[str] = None) -> 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")]
+    if extension:
+        udfs = [f for f in udfs if f.endswith(f".{extension}")]
+    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 user functions in default")
+
+    if CATALOG:
+        df = df.where(f"function == '{CATALOG}.default.udf_{udf}'")
+    else:
+        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).
+    """
+    if spark is None:
+        spark = SPARK
+    assert spark is not None
+
+    if not is_registered(udf, spark):
+        DEFAULT_LOGGER.debug(f"register udf {udf}")
+
+        if extension is None:
+            extension = get_extension(udf)
+
+        assert extension
+
+        path = PATH_UDFS.joinpath(f"{udf}.{extension}")
+
+        if extension == "sql":
+            spark.sql(path.get_sql())
+
+        elif extension == "py":
+            if not IS_UNITY_CATALOG:
+                assert path.exists(), f"udf not found ({path.string})"
+            else:
+                DEFAULT_LOGGER.debug(f"could not check if udf exists ({path.string})")
+
+            spec = importlib.util.spec_from_file_location(udf, path.string)
+            assert spec, f"no valid udf found ({path.string})"
+            assert spec.loader is not None
+
+            mod = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(mod)
+
+            u = UDFS[udf]
+            u(spark)
+
+        else:
+            raise ValueError(f"{udf} not found")
+
+
+def udf(name: str):
+    def decorator(fn: Callable):
+        UDFS[name] = fn
+        return fn
+
+    return decorator

fabricks/core/views.py

@@ -0,0 +1,41 @@
+from fabricks.context import PATH_VIEWS, SPARK
+from fabricks.context.log import DEFAULT_LOGGER
+from fabricks.utils.path import Path
+from fabricks.utils.sqlglot import fix as fix_sql
+
+
+def create_or_replace_view_internal(path: Path):
+    sql = path.get_sql()
+    file_name = path.get_file_name().split(".")[0]
+
+    try:
+        sql = f"""
+        create or replace view fabricks.{file_name}
+        as
+        {sql}
+        """
+        sql = fix_sql(sql)
+        DEFAULT_LOGGER.debug("create or replace (custom) view", extra={"label": f"fabricks.{file_name}", "sql": sql})
+
+        SPARK.sql(sql)
+
+    except Exception as e:
+        DEFAULT_LOGGER.exception(
+            "could not create nor replace (custom) view", extra={"label": f"fabricks.{file_name}", "exc_info": e}
+        )
+        raise e
+
+
+def create_or_replace_view(name: str):
+    p = PATH_VIEWS.joinpath(f"{name}.sql")
+    create_or_replace_view_internal(p)
+
+
+def create_or_replace_views():
+    DEFAULT_LOGGER.info("create or replace (custom) views")
+
+    for p in PATH_VIEWS.walk(file_format="sql", convert=True):
+        try:
+            create_or_replace_view_internal(p)
+        except Exception:
+            pass

fabricks/deploy/__init__.py

@@ -0,0 +1,92 @@
+import logging
+from typing import List, Optional, Union, cast
+
+from fabricks.context import FABRICKS_STORAGE
+from fabricks.context.log import DEFAULT_LOGGER
+from fabricks.core.jobs.base._types import Steps, TStep
+from fabricks.core.steps.base import BaseStep
+from fabricks.deploy.masks import deploy_masks
+from fabricks.deploy.notebooks import deploy_notebooks
+from fabricks.deploy.schedules import deploy_schedules
+from fabricks.deploy.tables import deploy_tables
+from fabricks.deploy.udfs import deploy_udfs
+from fabricks.deploy.utils import print_atomic_bomb
+from fabricks.deploy.views import deploy_views
+from fabricks.metastore.database import Database
+
+
+class Deploy:
+    @staticmethod
+    def tables(drop: bool = False):
+        deploy_tables(drop=drop)
+
+    @staticmethod
+    def views():
+        deploy_views()
+
+    @staticmethod
+    def udfs():
+        deploy_udfs()
+
+    @staticmethod
+    def masks():
+        deploy_masks()
+
+    @staticmethod
+    def notebooks():
+        deploy_notebooks()
+
+    @staticmethod
+    def schedules():
+        deploy_schedules()
+
+    @staticmethod
+    def armageddon(steps: Optional[Union[TStep, List[TStep], str, List[str]]], nowait: bool = False):
+        DEFAULT_LOGGER.warning("!💥 armageddon 💥!")
+        print_atomic_bomb(nowait=nowait)
+
+        DEFAULT_LOGGER.setLevel(logging.INFO)
+
+        if steps is None:
+            steps = Steps
+        assert steps is not None
+
+        if isinstance(steps, str):
+            steps = [cast(TStep, steps)]
+        elif isinstance(steps, List):
+            steps = [cast(TStep, s) for s in steps]
+        elif isinstance(steps, TStep):
+            steps = [steps]
+
+        fabricks = Database("fabricks")
+        fabricks.drop()
+
+        for s in steps:
+            step = BaseStep(s)
+            step.drop()
+
+        tmp = FABRICKS_STORAGE.joinpath("tmp")
+        tmp.rm()
+
+        checkpoint = FABRICKS_STORAGE.joinpath("checkpoints")
+        checkpoint.rm()
+
+        schema = FABRICKS_STORAGE.joinpath("schemas")
+        schema.rm()
+
+        schedule = FABRICKS_STORAGE.joinpath("schedules")
+        schedule.rm()
+
+        fabricks.create()
+
+        Deploy.tables(drop=True)
+        Deploy.udfs()
+        Deploy.masks()
+        Deploy.notebooks()
+
+        for s in steps:
+            step = BaseStep(s)
+            step.create()
+
+        Deploy.views()
+        Deploy.schedules()

fabricks/deploy/masks.py

@@ -0,0 +1,8 @@
+from fabricks.context.log import DEFAULT_LOGGER
+from fabricks.core.masks import register_all_masks
+
+
+def deploy_masks():
+    DEFAULT_LOGGER.info("create or replace masks")
+
+    register_all_masks()

fabricks/deploy/notebooks.py

@@ -0,0 +1,71 @@
+import base64
+import io
+import os
+from importlib import resources
+
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.service import workspace
+
+from fabricks.context import PATH_NOTEBOOKS
+from fabricks.context.log import DEFAULT_LOGGER
+
+
+def deploy_notebook(notebook: str):
+    from fabricks.api import notebooks
+
+    DEFAULT_LOGGER.debug(f"overwrite {notebook}")
+
+    w = WorkspaceClient()
+
+    target = f"{PATH_NOTEBOOKS}/{notebook}.py"
+    src = resources.files(notebooks) / f"{notebook}.py"
+
+    with io.open(src, "rb") as file:  # type: ignore
+        content = file.read()
+
+    encoded = base64.b64encode(content).decode("utf-8")
+
+    w.workspace.import_(
+        path=target,
+        content=encoded,
+        format=workspace.ImportFormat.AUTO,
+        language=workspace.Language.PYTHON,
+        overwrite=True,
+    )
+
+
+def deploy_notebooks():
+    DEFAULT_LOGGER.info("overwrite notebooks")
+
+    _create_dir_if_not_exists()
+    _clean_dir()
+
+    for n in [
+        "cluster",
+        "initialize",
+        "process",
+        "schedule",
+        "run",
+        "terminate",
+    ]:
+        deploy_notebook(notebook=n)
+
+
+def _create_dir_if_not_exists():
+    dir = str(PATH_NOTEBOOKS)
+    os.makedirs(dir, exist_ok=True)
+
+
+def _clean_dir():
+    dir = str(PATH_NOTEBOOKS)
+    for n in [
+        "cluster",
+        "initialize",
+        "process",
+        "schedule",
+        "run",
+        "terminate",
+    ]:
+        file_path = os.path.join(dir, f"{n}.py")
+        if os.path.isfile(file_path):
+            os.remove(file_path)

fabricks/deploy/schedules.py

@@ -0,0 +1,10 @@
+from fabricks.context.log import DEFAULT_LOGGER
+from fabricks.core.schedules import create_or_replace_views
+from fabricks.core.views import create_or_replace_views as create_or_replace_custom_views
+
+
+def deploy_schedules():
+    DEFAULT_LOGGER.info("create or replace schedules")
+
+    create_or_replace_custom_views()
+    create_or_replace_views()

fabricks/deploy/tables.py

@@ -0,0 +1,82 @@
+from pyspark.sql.types import LongType, StringType, StructField, StructType, TimestampType
+
+from fabricks.cdc import NoCDC
+from fabricks.context import SPARK
+from fabricks.context.log import DEFAULT_LOGGER
+from fabricks.metastore.table import Table
+
+
+def deploy_tables(drop: bool = False):
+    DEFAULT_LOGGER.info("create or replace fabricks (default) tables")
+
+    create_table_log(drop)
+    create_table_dummy(drop)
+    create_table_step(drop)
+
+
+def create_table_step(drop: bool = False):
+    table = Table("fabricks", "steps")
+    if drop:
+        table.drop()
+
+    if not table.exists():
+        schema = StructType(
+            [
+                StructField("step", StringType(), True),
+                StructField("expand", StringType(), True),
+                StructField("order", LongType(), True),
+            ]
+        )
+        table.create(schema=schema, partitioning=True, partition_by=["expand"])
+
+
+def create_table_log(drop: bool = False):
+    table = Table("fabricks", "logs")
+    if drop:
+        table.drop()
+
+    if not table.exists():
+        schema = StructType(
+            [
+                StructField("schedule_id", StringType(), True),
+                StructField("schedule", StringType(), True),
+                StructField("step", StringType(), True),
+                StructField("job_id", StringType(), True),
+                StructField("job", StringType(), True),
+                StructField("notebook_id", StringType(), True),
+                StructField("level", StringType(), True),
+                StructField("status", StringType(), True),
+                StructField("timestamp", TimestampType(), True),
+                StructField(
+                    "exception",
+                    StructType(
+                        [
+                            StructField("type", StringType(), True),
+                            StructField("message", StringType(), True),
+                            StructField("traceback", StringType(), True),
+                        ]
+                    ),
+                    True,
+                ),
+            ]
+        )
+        table.create(schema=schema, partitioning=True, partition_by=["schedule_id", "step"])
+
+
+def create_table_dummy(drop: bool = False):
+    cdc = NoCDC("fabricks", "dummy")
+
+    if drop:
+        cdc.drop()
+
+    if not cdc.table.exists():
+        df = SPARK.sql(
+            """
+            select
+            1 as __key,
+            md5('1') as __hash,
+            cast('1900-01-01' as timestamp) as __valid_from,
+            cast('9999-12-31' as timestamp) as __valid_to
+            """
+        )
+        cdc.overwrite(df)

fabricks/deploy/udfs.py

@@ -0,0 +1,19 @@
+from fabricks.context import SPARK
+from fabricks.context.log import DEFAULT_LOGGER
+from fabricks.core.udfs import register_all_udfs
+from fabricks.utils.sqlglot import fix as fix_sql
+
+
+def deploy_udfs():
+    DEFAULT_LOGGER.info("create or replace udfs")
+
+    register_all_udfs(extension="sql")
+    create_or_replace_udf_job_id()
+
+
+def create_or_replace_udf_job_id():
+    sql = "create or replace function fabricks.udf_job_id(job string) returns string return md5(job)"
+    sql = fix_sql(sql)
+
+    DEFAULT_LOGGER.debug("create or replace fabricks.udf_job_id", extra={"sql": sql})
+    SPARK.sql(sql)

fabricks/deploy/utils.py

@@ -0,0 +1,36 @@
+import time
+
+
+def print_atomic_bomb(nowait: bool = False):
+    def print_and_wait(message: str):
+        if not nowait:
+            time.sleep(0.5)
+        print(message)
+
+    print("")
+    print_and_wait("        ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡤⠤⠴⠾⠋⠉⠛⢾⡏⠙⠿⠦⠤⢤⣀⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀                    ")
+    print_and_wait("        ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡤⢶⣿⠉⢀⣀⡠⠆⠀⠀⠀⠀⠀⠀⠀⢤⣀⣀⠈⢹⣦⢤⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀                    ")
+    print_and_wait("        ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣴⣿⠁⢋⡙⠁⠀⡝⠀⠀⠀⠀⣀⡸⠋⠁⠀⠀⠹⡀⠀⠈⠈⠆⢹⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀                   ")
+    print_and_wait("        ⠀⠀⠀⠀⠀⠀⠀⢀⣠⣤⣿⣁⡡⣴⡏⠀⠀⠀⢀⠀⢧⣀⠄⠀⠀⠀⣀⣰⠆⢀⠁⠀⠀⢈⣶⡤⣀⢹⣦⣄⡀⠀⠀⠀⠀⠀⠀⠀⠀                  ")
+    print_and_wait("        ⠀⠀⠀⠀⠀⣠⢴⠟⢁⡝⠀⠁⠀⠃⠉⠀⠀⠘⣯⠀⡀⠾⣤⣄⣠⢤⠾⠄⠀⣸⠖⠀⠀⠈⠀⠃⠀⠀⠹⡄⠙⣶⢤⡀⠀⠀⠀⠀⠀                 ")
+    print_and_wait("        ⠀⠀⠀⣠⠾⡇⠈⣀⡞⠀⠀⠀⠀⡀⠀⢀⣠⣄⣇⠀⣳⠴⠃⠀⠀⠀⠣⢴⠉⣰⣇⣀⣀⠀⠀⡄⠀⠀⠀⢹⣄⡘⠈⡷⣦⠀⠀⠀⠀                 ")
+    print_and_wait("        ⢠⠞⠉⢻⡄⠀⠀⠈⠙⠀⠀⠀⠀⠙⣶⣏⣤⣤⠟⠉⠁⠀⠀⠀⠀⠀⠀⠀⠉⠙⢦⣱⣌⣷⠊⠀⠀⠀⠀⠈⠁⠀⠀⠀⡝⠉⠻⣄⠀                 ")
+    print_and_wait("        ⠛⢀⡠⢼⡇⠀⠀⢀⡄⠀⢀⣀⡽⠚⠁⠀⠀⠀⢠⡀⢠⣀⠠⣔⢁⡀⠀⣄⠀⡄⠀⠀⠀⠈⠑⠺⣄⡀⠀⠠⡀⠀⠀⢠⡧⠄⠀⠘⢧                ")
+    print_and_wait("        ⡶⠋⠀⠀⠈⣠⣈⣩⠗⠒⠋⠀⠀⠀⠀⣀⣠⣆⡼⣷⣞⠛⠻⡉⠉⡟⠒⡛⣶⠧⣀⣀⣀⠀⠀⠀⠀⠈⠓⠺⢏⣉⣠⠋⠀⠀⠀⢢⣸                ")
+    print_and_wait("        ⠇⠐⠤⠤⠖⠁⣿⣀⣀⠀⠀⠀⠀⠀⠉⠁⠈⠉⠙⠛⢿⣷⡄⢣⡼⠀⣾⣿⠧⠒⠓⠚⠛⠉⠀⠀⠀⠀⠀⢀⣀⣾⡉⠓⠤⡤⠄⠸⢿                ")
+    print_and_wait("        ⣆⣤⠀⠀⠠⠀⠈⠓⠈⠓⠤⡀⠀⠀⠀⠀⠀⠀⠀⠀⠈⣿⣿⢸⠀⢸⣿⠇⠀⠀⠀⠀⠀⠀⠀⠀⢀⡤⠒⠁⠰⠃⠀⠠⠀⠀⢀⣀⠞                  ")
+    print_and_wait("        ⠀⠉⠓⢲⣄⡈⢀⣠⠀⠀⠀⡸⠶⠂⠀⠀⢀⠀⠀⠤⠞⢻⡇⠀⠀⢘⡟⠑⠤⠄⠀⢀⠀⠀⠐⠲⢿⡀⠀⠀⢤⣀⢈⣀⡴⠖⠋⠀⠀                 ")
+    print_and_wait("        ⠀⠀⠀⠀⠈⠉⠉⠙⠓⠒⣾⣁⣀⣴⠀⣀⠙⢧⠂⢀⣆⣀⣷⣤⣀⣾⣇⣀⡆⠀⢢⠛⢁⠀⢰⣀⣀⣹⠒⠒⠛⠉⠉⠉⠀⠀⠀⠀⠀                ")
+    print_and_wait("        ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠉⠁⠈⠉⠉⠛⠉⠙⠉⠀⠀⣿⡟⣿⣿⠀⠀⠈⠉⠉⠙⠋⠉⠉⠀⠉⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀                   ")
+    print_and_wait("        ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣿⡇⢻⣿⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀                      ")
+    print_and_wait("        ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⣤⣶⣾⣿⣿⠁⠀⢹⡛⣟⡶⢤⣀⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀                    ")
+    print_and_wait("        ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣴⠛⢯⣽⡟⢿⣿⠛⠿⠳⠞⠻⣿⠻⣆⢽⠟⣶⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀                   ")
+    print_and_wait("        ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠛⠃⠲⠯⠴⣦⣼⣷⣤⣤⣶⣤⣩⡧⠽⠷⠐⠛⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀                   ")
+    print_and_wait("        ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣿⡇⠀⣿⡆⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀                     ")
+    print_and_wait("        ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣀⣄⡀⢀⣀⣠⡾⡿⢡⢐⠻⣿⣄⣀⡀⠀⣀⣄⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀                  ")
+    print_and_wait("        ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣤⢴⡏⠁⠀⠝⠉⣡⠟⣰⠃⢸⣿⠀⣷⠙⢧⡉⠻⡅⠀⠙⡷⢤⣀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀                 ")
+    print_and_wait("        ⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣴⡟⠀⠈⣿⢄⡴⠞⠻⣄⣰⣡⠤⣞⣸⡤⢬⣧⣀⡿⠛⠦⣤⣶⡃⠀⢹⣦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀                 ")
+    print_and_wait("        ⠀⠀⠀⠀⠀⠀⢀⣴⣶⡿⠃⠉⢺⠁⠙⠒⠀⠀⣠⡉⠀⠉⠚⠉⠉⠑⠈⠀⠈⣧⠀⠀⠒⠋⠀⡹⠋⠀⢻⡶⠶⡄⠀⠀⠀⠀⠀⠀⠀                 ")
+    print_and_wait("        ⠀⠀⠀⠀⠀⣠⣾⣿⣇⠁⢈⡦⠀⡍⠋⠁⡀⠸⡋⠀⠀⠀⢘⠏⠉⡏⠀⠀⠀⢉⡷⠀⡌⠉⠋⡇⠠⣏⠈⢁⣦⣿⣦⠀⠀⠀⠀⠀⠀                 ")
+    print_and_wait("        ⠀⠀⠀⠀⠀⠉⣁⠀⠉⠉⠉⠙⠛⠛⠒⠚⠳⠤⢼⣤⣠⠤⣮⣠⣤⣼⠦⢤⣤⣿⠤⠾⠓⠒⠛⢓⠛⠉⠉⠉⠀⠈⠉⠀⠀⠀⠀⠀⠀                ")
+    print("")