fabricks 2024.7.1.5__py3-none-any.whl

fabricks/context/runtime.py

@@ -0,0 +1,143 @@
+import os
+import sys
+from typing import Final, List
+
+import yaml
+from databricks.sdk.runtime import spark
+
+from fabricks.utils.path import Path
+
+try:
+    runtime = Path(os.environ["FABRICKS_RUNTIME"], assume_git=True)
+    assert runtime, "runtime mandatory in cluster config"
+    PATH_RUNTIME: Final[Path] = runtime
+
+    notebooks = Path(os.environ["FABRICKS_NOTEBOOKS"], assume_git=True)
+    assert notebooks, "notebooks mandatory in cluster config"
+    PATH_NOTEBOOKS: Final[Path] = notebooks
+
+    version = os.environ["FABRICKS_VERSION"]
+    assert version, "version mandatory in cluster config"
+    VERSION: Final[str] = version
+
+    PATH_LIBRARIES = "/dbfs/mnt/fabricks/site-packages"
+    spark._sc._python_includes.append(PATH_LIBRARIES)  # type: ignore
+    sys.path.append(PATH_LIBRARIES)
+
+    try:
+        is_test = os.environ["FABRICKS_IS_TEST"] == "TRUE"
+    except Exception:
+        is_test = False
+    IS_TEST: Final[bool] = is_test
+
+    try:
+        is_debug = os.environ["FABRICKS_IS_DEBUG"] == "TRUE"
+    except Exception:
+        is_debug = False
+    IS_DEBUG: Final[bool] = is_debug
+
+    try:
+        is_live = os.environ["FABRICKS_IS_LIVE"] == "TRUE"
+    except Exception:
+        is_live = False
+    IS_LIVE: Final[bool] = is_live
+
+    conf_path = PATH_RUNTIME.join(
+        "fabricks",
+        f"conf.{spark.conf.get('spark.databricks.clusterUsageTags.clusterOwnerOrgId')}.yml",
+    )
+    with open(conf_path.string) as f:
+        data = yaml.safe_load(f)
+
+    conf: dict = [d["conf"] for d in data][0]
+    assert conf, "conf mandatory"
+    CONF_RUNTIME: Final[dict] = conf
+
+    BRONZE = CONF_RUNTIME.get("bronze", [{}])
+    SILVER = CONF_RUNTIME.get("silver", [{}])
+    GOLD = CONF_RUNTIME.get("gold", [{}])
+    STEPS = BRONZE + SILVER + GOLD
+
+    databases = CONF_RUNTIME.get("databases", [{}])
+    credentials = CONF_RUNTIME.get("credentials", {})
+    variables = CONF_RUNTIME.get("variables", {})
+    VARIABLES: dict = variables
+
+    conf_options = CONF_RUNTIME.get("options", {})
+    assert conf_options, "options mandatory"
+
+    secret_scope = conf_options.get("secret_scope")
+    assert secret_scope, "secret_scope mandatory in options"
+    SECRET_SCOPE: Final[str] = secret_scope
+
+    path_options = CONF_RUNTIME.get("path_options", {})
+    assert path_options, "options mandatory"
+
+    fabricks_uri = path_options.get("storage")
+    assert fabricks_uri, "storage mandatory in path options"
+    FABRICKS_STORAGE: Final[Path] = Path.from_uri(fabricks_uri, regex=variables)
+
+    path_udfs = path_options.get("udfs")
+    assert path_udfs, "udfs mandatory in path options"
+    PATH_UDFS: Final[Path] = PATH_RUNTIME.join(path_udfs)
+
+    path_parsers = path_options.get("parsers")
+    assert path_parsers, "parsers mandatory in path options"
+    PATH_PARSERS: Final[Path] = PATH_RUNTIME.join(path_parsers)
+
+    path_extenders = path_options.get("extenders")
+    assert path_extenders, "extenders mandatory in path options"
+    PATH_EXTENDERS: Final[Path] = PATH_RUNTIME.join(path_extenders)
+
+    path_views = path_options.get("views")
+    assert path_views, "views mandatory in path options"
+    PATH_VIEWS: Final[Path] = PATH_RUNTIME.join(path_views)
+
+    path_schedules = path_options.get("schedules")
+    assert path_schedules, "schedules mandatory in path options"
+    PATH_SCHEDULES: Final[Path] = PATH_RUNTIME.join(path_schedules)
+
+    path_requirements = path_options.get("requirements")
+    assert path_requirements, "requirements mandatory in path options"
+    PATH_REQUIREMENTS: Final[Path] = PATH_RUNTIME.join(path_requirements)
+
+    def _get_storage_paths(objects: List[dict]) -> dict:
+        d = {}
+        for o in objects:
+            if o:
+                name = o.get("name")
+                assert name
+                uri = o.get("path_options", {}).get("storage")
+                assert uri
+                d[name] = Path.from_uri(uri, regex=variables)
+        return d
+
+    PATHS_STORAGE: Final[dict[str, Path]] = {
+        "fabricks": FABRICKS_STORAGE,
+        **_get_storage_paths(BRONZE),
+        **_get_storage_paths(SILVER),
+        **_get_storage_paths(GOLD),
+        **_get_storage_paths(databases),
+    }
+
+    def _get_runtime_path(objects: List[dict]) -> dict:
+        d = {}
+        for o in objects:
+            name = o.get("name")
+            assert name
+            uri = o.get("path_options", {}).get("runtime")
+            assert uri
+            d[name] = PATH_RUNTIME.join(uri)
+        return d
+
+    PATHS_RUNTIME: Final[dict[str, Path]] = {
+        **_get_runtime_path(BRONZE),
+        **_get_runtime_path(SILVER),
+        **_get_runtime_path(GOLD),
+    }
+
+except KeyError as e:
+    raise e
+
+except AssertionError as e:
+    raise e

fabricks/context/spark.py

@@ -0,0 +1,43 @@
+from typing import Optional, Tuple
+
+from pyspark.dbutils import DBUtils
+from pyspark.sql import SparkSession
+
+from fabricks.context.runtime import CONF_RUNTIME, SECRET_SCOPE
+from fabricks.utils.secret import add_secret_to_spark, get_secret_from_secret_scope
+
+
+def build_spark_session(new: Optional[bool] = False, log: Optional[bool] = False) -> Tuple[SparkSession, DBUtils]:
+    if new:
+        spark = SparkSession.builder.getOrCreate().newSession()  # type: ignore
+
+        catalog = CONF_RUNTIME.get("options", {}).get("catalog")
+        if catalog:
+            spark.sql(f"use catalog {catalog};")
+
+        # delta
+        spark.sql("set spark.databricks.delta.schema.autoMerge.enabled = True;")
+        spark.sql("set spark.databricks.delta.resolveMergeUpdateStructsByName.enabled = True;")
+
+        spark_options = CONF_RUNTIME.get("spark_options", {})
+        if spark_options:
+            sql_options = spark_options.get("sql", {})
+            for key, value in sql_options.items():
+                spark.sql(f"set {key} = {value};")
+
+            conf_options = spark_options.get("conf", {})
+            for key, value in conf_options.items():
+                spark.conf.set(key, value)
+
+        credentials = CONF_RUNTIME.get("credentials", {})
+        for uri, secret in credentials.items():
+            s = get_secret_from_secret_scope(secret_scope=SECRET_SCOPE, name=secret)
+            add_secret_to_spark(secret=s, uri=uri)
+
+    else:
+        spark = SparkSession.builder.getOrCreate()  # type: ignore
+
+    return spark, DBUtils(spark)
+
+
+build_spark_session(new=True, log=True)

fabricks/context/types.py

@@ -0,0 +1,123 @@
+from typing import List, Optional, TypedDict
+
+
+class RuntimePathOptions(TypedDict):
+    storage: str
+    udfs: str
+    extenders: str
+    parsers: str
+    schedules: str
+    views: str
+    requirements: str
+
+
+class RuntimeTimeoutOptions(TypedDict):
+    step: int
+    job: int
+    pre_run: int
+    post_run: int
+
+
+class StepTimeoutOptions(TypedDict):
+    step: Optional[int]
+    job: Optional[int]
+    pre_run: Optional[int]
+    post_run: Optional[int]
+
+
+class RuntimeOptions(TypedDict):
+    secret_scope: str
+    catalog: str
+    workers: int
+    timeouts: RuntimeTimeoutOptions
+    retention_days: int
+
+
+class SparkOptions(TypedDict):
+    sql: dict
+    conf: dict
+
+
+class StepPathOptions(TypedDict):
+    runtime: str
+    storage: str
+
+
+class InvokeOptions(TypedDict):
+    notebook: str
+    arguments: Optional[dict[str, str]]
+
+
+class StepOptions(TypedDict):
+    order: int
+    workers: Optional[int]
+    timeouts: StepTimeoutOptions
+    extender: Optional[str]
+    pre_run: Optional[InvokeOptions]
+    post_run: Optional[InvokeOptions]
+
+
+class SilverOptions(StepOptions):
+    parent: str
+    stream: Optional[bool]
+    local_checkpoint: Optional[bool]
+
+
+class GoldOptions(StepOptions):
+    schema_drift: Optional[bool]
+
+
+class Step(TypedDict):
+    name: str
+
+
+class TableOptions(TypedDict):
+    powerbi: Optional[bool]
+    liquid_clustering: Optional[bool]
+    properties: Optional[dict[str, str]]
+    retention_days: Optional[int]
+
+
+class Bronze(Step):
+    options: StepOptions
+    path_options: StepPathOptions
+    table_options: Optional[TableOptions]
+
+
+class Silver(Step):
+    options: SilverOptions
+    path_options: StepPathOptions
+    table_options: Optional[TableOptions]
+
+
+class Gold(Step):
+    options: GoldOptions
+    path_options: StepPathOptions
+    table_options: Optional[TableOptions]
+
+
+class PowerBI(Step):
+    pass
+
+
+class DatabasePathOptions(TypedDict):
+    storage: str
+
+
+class Database(TypedDict):
+    name: str
+    path_options: DatabasePathOptions
+
+
+class Conf(TypedDict):
+    name: str
+    options: RuntimeOptions
+    path_options: RuntimePathOptions
+    spark_options: SparkOptions
+    bronze: Optional[List[Bronze]]
+    silver: Optional[List[Silver]]
+    gold: Optional[List[Gold]]
+    powerbi: Optional[List[PowerBI]]
+    databases: Optional[List[Database]]
+    variables: Optional[List[dict[str, str]]]
+    credentials: Optional[List[dict[str, str]]]

fabricks/core/__init__.py

@@ -0,0 +1,4 @@
+from fabricks.core.jobs import Bronzes, Golds, Silvers, Steps, get_job, get_jobs
+from fabricks.core.steps import get_step
+
+__all__ = ["get_job", "get_jobs", "get_step", "Bronzes", "Silvers", "Golds", "Steps"]

fabricks/core/dags/__init__.py

@@ -0,0 +1,9 @@
+from fabricks.core.dags.generator import DagGenerator
+from fabricks.core.dags.processor import DagProcessor
+from fabricks.core.dags.terminator import DagTerminator
+
+__all__ = [
+    "DagGenerator",
+    "DagProcessor",
+    "DagTerminator",
+]

fabricks/core/dags/base.py

@@ -0,0 +1,72 @@
+import re
+from typing import Optional, cast
+
+from databricks.sdk.runtime import spark
+from pyspark.sql import DataFrame
+from pyspark.sql.functions import expr
+
+from fabricks.context import FABRICKS_STORAGE
+from fabricks.core.dags.log import DagsTableLogger
+from fabricks.metastore.table import Table
+from fabricks.utils.azure_table import AzureTable
+from fabricks.utils.secret import AccessKey, get_secret_from_secret_scope
+
+
+class BaseDags:
+    def __init__(self, schedule_id: str):
+        self.schedule_id = schedule_id
+
+    def get_connection_string(self) -> str:
+        storage_account = FABRICKS_STORAGE.get_storage_account()
+        secret = get_secret_from_secret_scope("bmskv", f"{storage_account}-access-key")
+        access_key = cast(AccessKey, secret).key
+        connection_string = f"DefaultEndpointsProtocol=https;AccountName={storage_account};AccountKey={access_key};EndpointSuffix=core.windows.net"
+        return connection_string
+
+    def get_table(self) -> AzureTable:
+        cs = self.get_connection_string()
+        table = AzureTable(f"t{self.schedule_id}", connection_string=cs)
+        return table
+
+    def get_logs(self, step: Optional[str] = None) -> DataFrame:
+        q = f"PartitionKey eq '{self.schedule_id}'"
+        if step:
+            q += f" and Step eq '{step}'"
+
+        d = DagsTableLogger.table.query(q)
+        df = spark.createDataFrame(d)
+        if "Exception" not in df.columns:
+            df = df.withColumn("Exception", expr("null"))
+
+        df = spark.sql(
+            """
+            select
+              ScheduleId as schedule_id,
+              Schedule as schedule,
+              Step as step,
+              JobId as job_id,
+              Job as job,
+              NotebookId as notebook_id,
+              `Level` as `level`,
+              `Message` as `status`,
+              to_timestamp(`Created`, 'dd/MM/yy HH:mm:ss') as `timestamp`,
+              from_json(Exception, 'type STRING, message STRING, traceback STRING') as exception
+            from
+              {df}
+            """,
+            df=df,
+        )
+        return df
+
+    def write_logs(self, df: DataFrame):
+        (
+            df.write.format("delta")
+            .mode("overwrite")
+            .option("mergeSchema", "true")
+            .option("partitionOverwriteMode", "dynamic")
+            .save(Table("fabricks", "logs").deltapath.string)
+        )
+
+    def remove_invalid_characters(self, s: str) -> str:
+        out = re.sub("[^a-zA-Z0-9]", "", s)
+        return out

fabricks/core/dags/generator.py

@@ -0,0 +1,154 @@
+import time
+from typing import Optional, Tuple
+from uuid import uuid4
+
+from databricks.sdk.runtime import spark
+from pyspark.sql import DataFrame
+
+from fabricks.core.dags.base import BaseDags
+from fabricks.core.dags.log import DagsTableLogger
+from fabricks.utils.azure_queue import AzureQueue
+
+
+class DagGenerator(BaseDags):
+    def __init__(self, schedule: str):
+        self.schedule = schedule
+        schedule_id = str(uuid4().hex)
+        super().__init__(schedule_id=schedule_id)
+
+    def get_jobs(self) -> DataFrame:
+        return spark.sql(
+            f"""
+            with logs as (
+              select 
+                l.job_id,
+                median(l.duration) as median_duration
+              from
+                fabricks.logs_pivot l
+              where
+                true
+                and duration is not null
+                and date_diff(day, l.start_time , current_date) < 10
+              group by
+                l.job_id
+            )
+            select
+              'statuses' as PartitionKey,
+              '{self.schedule_id}' as ScheduleId,
+              '{self.schedule}' as Schedule,
+              j.job_id::string as RowKey,
+              j.step as Step,
+              j.job_id as JobId,
+              j.job as Job,
+              'scheduled' as `Status`,
+              max(median_duration) as `MedianDuration`, 
+              dense_rank() over (order by max(median_duration) desc) as Rank
+            from
+              fabricks.jobs j
+              inner join fabricks.{self.schedule}_schedule v on j.job_id = v.job_id
+              left join logs l on j.job_id = l.job_id
+            group by all
+            """
+        )
+
+    def get_dependencies(self, job_df: Optional[DataFrame] = None) -> DataFrame:
+        if job_df is None:
+            job_df = self.get_jobs()
+
+        return spark.sql(
+            """
+            select
+              'dependencies' as PartitionKey, 
+              d.dependency_id::string as RowKey,
+              {schedule_id} as ScheduleId,
+              {schedule} as Schedule,
+              d.dependency_id as DependencyId,
+              j.Step as Step,
+              j.Job as Job,
+              j.JobId as JobId,
+              p.Step as ParentStep,
+              p.Job as Parent,
+              p.JobId as ParentId
+            from
+              fabricks.dependencies d
+              inner join {job} j on d.job_id = j.JobId
+              inner join {job} p on d.parent_id = p.JobId
+            where
+              true
+              and d.parent_id is not null
+              and not d.job_id = d.parent_id
+              and not exists (
+                select 1
+                from 
+                  fabricks.dependencies_circular dc
+                where
+                  true 
+                  and d.job_id = dc.job_id 
+                  and d.parent_id = dc.parent_id
+
+              )
+            group by all
+            """,
+            job=job_df,
+            schedule=self.schedule,
+            schedule_id=self.schedule_id,
+        )
+
+    def get_steps(self, job_df: Optional[DataFrame] = None) -> DataFrame:
+        if job_df is None:
+            job_df = self.get_jobs()
+
+        return spark.sql(
+            """
+            select
+              Step
+            from
+              {job}
+            group by
+              Step
+            """,
+            job=job_df,
+        )
+
+    def generate(self) -> Tuple[str, DataFrame, DataFrame]:
+        job_df = self.get_jobs()
+        deps_df = self.get_dependencies(job_df)
+        step_df = self.get_steps(job_df)
+
+        table = self.get_table()
+        table.create_if_not_exists()
+        table.truncate_all_partitions()
+        table.upsert(job_df)
+        table.upsert(deps_df)
+
+        df = spark.sql(
+            """
+            select
+              ScheduleId as PartitionKey,
+              ScheduleId,
+              `Schedule`,
+              Step,
+              Job,
+              JobId,
+              date_format(current_timestamp(), 'dd/MM/yy HH:mm:ss') as Created,
+              'INFO' as `Level`,
+              `Status` as `Message`,
+              from_json(null, 'type STRING, message STRING, traceback STRING') as Exception,
+              md5(array_join(array(ScheduleId, `Schedule`, Step, Job, JobId, Created,  `Level`, `Message`, -1), "*")) as RowKey
+            from
+              {df}        
+            """,
+            df=job_df,
+        )
+        DagsTableLogger.table.upsert(df)
+
+        cs = self.get_connection_string()
+        for row in step_df.collect():
+            step = self.remove_invalid_characters(row.Step)
+            queue = AzureQueue(f"q{step}{self.schedule_id}", connection_string=cs)
+            queue.create_if_not_exists()
+            queue.clear()
+
+        time.sleep(60)
+
+        return self.schedule_id, job_df, deps_df

fabricks/core/dags/log.py

@@ -0,0 +1,14 @@
+import logging
+from typing import cast
+
+from fabricks.context.runtime import FABRICKS_STORAGE, SECRET_SCOPE
+from fabricks.utils.azure_table import AzureTable
+from fabricks.utils.log import get_logger
+from fabricks.utils.secret import AccessKey, get_secret_from_secret_scope
+
+storage_account = FABRICKS_STORAGE.get_storage_account()
+secret = get_secret_from_secret_scope(SECRET_SCOPE, f"{storage_account}-access-key")
+access_key = cast(AccessKey, secret).key
+
+table = AzureTable("dags", storage_account=storage_account, access_key=access_key)
+DagsLogger, DagsTableLogger = get_logger("dags", logging.DEBUG, table=table)