fabricks 3.0.5.2__py3-none-any.whl → 3.0.7__py3-none-any.whl

fabricks/api/__init__.py

@@ -1,9 +1,11 @@
 from fabricks.api.context import init_spark_session
 from fabricks.api.core import get_job, get_jobs, get_step
+from fabricks.api.deploy import Deploy
 
 __all__ = [
     "init_spark_session",
     "get_job",
     "get_jobs",
     "get_step",
+    "Deploy",
 ]

fabricks/api/context.py

@@ -1,5 +1,4 @@
-from fabricks.context import BRONZE, DBUTILS, GOLD, SECRET_SCOPE, SILVER, SPARK, init_spark_session
-from fabricks.context.runtime import pprint_runtime
+from fabricks.context import BRONZE, DBUTILS, GOLD, SECRET_SCOPE, SILVER, SPARK, init_spark_session, pprint_runtime
 from fabricks.core.jobs.base._types import Bronzes, Golds, Silvers, Steps
 
 # step

fabricks/api/deploy.py

@@ -0,0 +1,3 @@
+from fabricks.deploy import Deploy
+
+__all__ = ["Deploy"]

fabricks/api/job_schema.py

@@ -1,3 +1,3 @@
-from fabricks.core.scripts.job_schema import get_job_schema
+from fabricks.core.job_schema import get_job_schema, print_job_schema
 
-__all__ = ["get_job_schema"]
+__all__ = ["get_job_schema", "print_job_schema"]

fabricks/api/masks.py

@@ -0,0 +1,3 @@
+from fabricks.core.masks import register_all_masks, register_mask
+
+__all__ = ["register_all_masks", "register_mask"]

fabricks/api/notebooks/initialize.py

@@ -1,11 +1,11 @@
 # Databricks notebook source
-# MAGIC %run ./add_fabricks
+# MAGIC %run ./add_missing_modules
 
 # COMMAND ----------
 
 from databricks.sdk.runtime import dbutils, display
 
-from fabricks.core.scripts import generate
+from fabricks.core.schedules import generate
 
 # COMMAND ----------
 

fabricks/api/notebooks/process.py

@@ -1,12 +1,12 @@
 # Databricks notebook source
-# MAGIC %run ./add_fabricks
+# MAGIC %run ./add_missing_modules
 
 # COMMAND ----------
 
 from databricks.sdk.runtime import dbutils
 from pyspark.errors.exceptions.base import IllegalArgumentException
 
-from fabricks.core.scripts import process
+from fabricks.core.schedules import process
 
 # COMMAND ----------
 

fabricks/api/notebooks/run.py

@@ -1,5 +1,5 @@
 # Databricks notebook source
-# MAGIC %run ./add_fabricks
+# MAGIC %run ./add_missing_modules
 
 # COMMAND ----------
 
@@ -7,7 +7,7 @@ import json
 
 from databricks.sdk.runtime import dbutils
 
-from fabricks.core.dags.run import run
+from fabricks.core.schedules import run
 
 # COMMAND ----------
 

fabricks/api/notebooks/schedule.py

@@ -0,0 +1,75 @@
+# Databricks notebook source
+# MAGIC %run ./add_missing_modules
+
+# COMMAND ----------
+
+from logging import DEBUG
+from typing import Any, cast
+
+from databricks.sdk.runtime import dbutils, display, spark
+
+from fabricks.context import PATH_NOTEBOOKS
+from fabricks.context.log import DEFAULT_LOGGER
+from fabricks.core import get_step
+from fabricks.core.jobs.base._types import TStep
+from fabricks.core.schedules import generate, terminate
+from fabricks.utils.helpers import run_in_parallel, run_notebook
+
+# COMMAND ----------
+
+DEFAULT_LOGGER.setLevel(DEBUG)
+
+# COMMAND ----------
+
+dbutils.widgets.text("schedule", "---")
+
+# COMMAND ----------
+
+schedule = dbutils.widgets.get("schedule")
+assert schedule != "---", "no schedule provided"
+
+# COMMAND ----------
+
+schedule_id, job_df, dependency_df = generate(schedule=schedule)
+
+# COMMAND ----------
+
+print(schedule_id)
+
+# COMMAND ----------
+
+display(job_df)
+
+# COMMAND ----------
+
+display(dependency_df)
+
+# COMMAND ----------
+steps = [row.step for row in spark.sql("select step from {df} group by step", df=job_df).collect()]
+
+# COMMAND ----------
+
+
+def _schedule(task: Any):
+    step = get_step(step=cast(TStep, task))
+    run_notebook(
+        PATH_NOTEBOOKS.joinpath("process"),
+        timeout=step.timeouts.step,
+        step=task,
+        schedule_id=schedule_id,
+        schedule=schedule,
+        workers=step.workers,
+    )
+
+
+# COMMAND ----------
+
+run_in_parallel(_schedule, steps)
+
+# COMMAND ----------
+
+terminate(schedule_id=schedule_id)
+
+# COMMAND ----------
+
+dbutils.notebook.exit(value="exit (0)")  # type: ignore

fabricks/api/notebooks/terminate.py

@@ -1,12 +1,12 @@
 # Databricks notebook source
-# MAGIC %run ./add_fabricks
+# MAGIC %run ./add_missing_modules
 
 # COMMAND ----------
 
 from databricks.sdk.runtime import dbutils
 from pyspark.errors.exceptions.base import IllegalArgumentException
 
-from fabricks.core.scripts import terminate
+from fabricks.core.schedules import terminate
 
 # COMMAND ----------
 

fabricks/api/schedules.py

@@ -1,17 +1,3 @@
-from fabricks.core.schedules import (
-    create_or_replace_view,
-    create_or_replace_views,
-    get_dependencies,
-    get_mermaid_diagram,
-    get_schedule,
-    get_schedules,
-)
+from fabricks.core.schedules import create_or_replace_view, create_or_replace_views, generate, process, terminate
 
-__all__ = [
-    "create_or_replace_view",
-    "create_or_replace_views",
-    "get_dependencies",
-    "get_mermaid_diagram",
-    "get_schedule",
-    "get_schedules",
-]
+__all__ = ["create_or_replace_view", "create_or_replace_views", "terminate", "generate", "process"]

fabricks/cdc/__init__.py

@@ -1,4 +1,4 @@
-from fabricks.cdc.base import BaseCDC, ChangeDataCaptures
+from fabricks.cdc.base import AllowedChangeDataCaptures, BaseCDC
 from fabricks.cdc.cdc import CDC
 from fabricks.cdc.nocdc import NoCDC
 from fabricks.cdc.scd1 import SCD1
@@ -7,7 +7,7 @@ from fabricks.cdc.scd2 import SCD2
 __all__ = [
     "BaseCDC",
     "CDC",
-    "ChangeDataCaptures",
+    "AllowedChangeDataCaptures",
     "NoCDC",
     "SCD1",
     "SCD2",

fabricks/cdc/base/__init__.py

@@ -1,4 +1,4 @@
-from fabricks.cdc.base._types import ChangeDataCaptures
+from fabricks.cdc.base._types import AllowedChangeDataCaptures
 from fabricks.cdc.base.cdc import BaseCDC
 
-__all__ = ["BaseCDC", "ChangeDataCaptures"]
+__all__ = ["BaseCDC", "AllowedChangeDataCaptures"]

fabricks/cdc/base/_types.py

@@ -1,3 +1,10 @@
-from typing import Literal
+from __future__ import annotations
 
-ChangeDataCaptures = Literal["nocdc", "scd1", "scd2"]
+from typing import Literal, Union
+
+from pyspark.sql import DataFrame
+
+from fabricks.metastore.table import Table
+
+AllowedChangeDataCaptures = Literal["nocdc", "scd1", "scd2"]
+AllowedSources = Union[DataFrame, Table, str]

fabricks/cdc/base/configurator.py

@@ -4,11 +4,13 @@ from abc import ABC, abstractmethod
 from typing import List, Optional, Union
 
 from pyspark.sql import DataFrame, SparkSession
-from pyspark.sql.connect.dataframe import DataFrame as CDataFrame
 
+from fabricks.cdc.base._types import AllowedSources
 from fabricks.context import SPARK
+from fabricks.context.log import DEFAULT_LOGGER
 from fabricks.metastore.database import Database
 from fabricks.metastore.table import Table
+from fabricks.utils._types import DataFrameLike
 
 
 class Configurator(ABC):
@@ -34,25 +36,23 @@ class Configurator(ABC):
         return self.table.is_view
 
     @property
-    def is_registered(self):
-        return self.table.is_registered
+    def registered(self):
+        return self.table.registered
 
     @property
     def qualified_name(self):
         return f"{self.database}_{'_'.join(self.levels)}"
 
     @abstractmethod
-    def get_query(self, src: Union[DataFrame, Table, str], **kwargs):
-        raise NotImplementedError()
+    def get_query(self, src: AllowedSources, **kwargs) -> str: ...
 
     @abstractmethod
-    def get_data(self, src: Union[DataFrame, Table, str], **kwargs) -> DataFrame:
-        raise NotImplementedError()
+    def get_data(self, src: AllowedSources, **kwargs) -> DataFrame: ...
 
     @abstractmethod
     def create_table(
         self,
-        src: Union[DataFrame, Table, str],
+        src: AllowedSources,
         partitioning: Optional[bool] = False,
         partition_by: Optional[Union[List[str], str]] = None,
         identity: Optional[bool] = False,
@@ -60,19 +60,32 @@ class Configurator(ABC):
         cluster_by: Optional[Union[List[str], str]] = None,
         properties: Optional[dict[str, str]] = None,
         **kwargs,
-    ):
-        raise NotImplementedError()
+    ): ...
 
     @abstractmethod
-    def drop(self):
-        raise NotImplementedError()
+    def drop(self): ...
 
     @abstractmethod
-    def create_or_replace_view(self, src: Union[Table, str], **kwargs):
-        raise NotImplementedError()
+    def create_or_replace_view(self, src: Union[Table, str], **kwargs): ...
 
     @property
-    def allowed_leading_columns(self):
+    def allowed_input__columns(self) -> List[str]:
+        cols = self.__columns
+
+        if self.slowly_changing_dimension:
+            if "__valid_from" in cols:
+                cols.remove("__valid_from")
+            if "__valid_to" in cols:
+                cols.remove("__valid_to")
+            if "__is_current" in cols:
+                cols.remove("__is_current")
+            if "__is_deleted" in cols:
+                cols.remove("__is_deleted")
+
+        return cols
+
+    @property
+    def allowed_ouput_leading__columns(self) -> List[str]:
         cols = [
             "__identity",
             "__source",
@@ -93,7 +106,7 @@ class Configurator(ABC):
         return cols
 
     @property
-    def allowed_trailing_columns(self):
+    def allowed_output_trailing__columns(self) -> List[str]:
         cols = [
             "__operation",
             "__metadata",
@@ -101,19 +114,36 @@ class Configurator(ABC):
             "__rescued_data",
         ]
 
-        if self.change_data_capture == "scd1":
-            cols.remove("__operation")
-        elif self.change_data_capture == "scd2":
+        if self.slowly_changing_dimension:
             cols.remove("__operation")
 
         return cols
 
+    @property
+    def __columns(self) -> List[str]:
+        return [
+            # Leading
+            "__identity",
+            "__source",
+            "__key",
+            "__timestamp",
+            "__valid_from",
+            "__valid_to",
+            "__is_current",
+            "__is_deleted",
+            # Trailing
+            "__operation",
+            "__metadata",
+            "__hash",
+            "__rescued_data",
+        ]
+
     @property
     def slowly_changing_dimension(self) -> bool:
         return self.change_data_capture in ["scd1", "scd2"]
 
-    def get_src(self, src: Union[DataFrame, Table, str]) -> DataFrame:
-        if isinstance(src, (DataFrame, CDataFrame)):
+    def get_src(self, src: AllowedSources) -> DataFrame:
+        if isinstance(src, DataFrameLike):
             df = src
         elif isinstance(src, Table):
             df = self.table.dataframe
@@ -124,55 +154,70 @@ class Configurator(ABC):
 
         return df
 
-    def has_data(self, src: Union[DataFrame, Table, str], **kwargs) -> bool:
+    def has_data(self, src: AllowedSources, **kwargs) -> bool:
+        DEFAULT_LOGGER.debug("check if has data", extra={"label": self})
         df = self.get_src(src=src)
         return not df.isEmpty()
 
-    def get_columns(self, src: Union[DataFrame, Table, str], backtick: Optional[bool] = True) -> List[str]:
+    def get_columns(
+        self,
+        src: AllowedSources,
+        backtick: Optional[bool] = True,
+        sort: Optional[bool] = True,
+        check: Optional[bool] = True,
+    ) -> List[str]:
         if backtick:
             backtick = True
 
         df = self.get_src(src=src)
         columns = df.columns
 
+        if check:
+            for c in columns:
+                # avoid duplicate column issue in merge
+                if c.startswith("__") and c in self.__columns:
+                    assert c in self.allowed_input__columns, f"{c} is not allowed"
+
+        if sort:
+            columns = self.sort_columns(columns)
+
         if backtick:
             return [f"`{c}`" for c in columns]
         else:
             return columns
 
-    def reorder_columns(self, df: DataFrame) -> DataFrame:
-        fields = [f"`{c}`" for c in df.columns if not c.startswith("__")]
+    def sort_columns(self, columns: List[str]) -> List[str]:
+        fields = [c for c in columns if not c.startswith("__")]
+
+        leading = self.allowed_ouput_leading__columns
+        trailing = self.allowed_output_trailing__columns
 
-        leading = self.allowed_leading_columns
-        trailing = self.allowed_trailing_columns
-        if (
-            "__key" not in df.columns and "__hash" in df.columns
-        ):  # move __hash to the front of the table to ensure statistics are present
+        # move __hash to the front of the table to ensure statistics are present
+        if "__key" not in columns and "__hash" in columns:
             leading = ["__hash" if c == "__key" else c for c in leading]
             trailing = [c for c in trailing if c != "__hash"]
 
-        __leading = [c for c in leading if c in df.columns]
-        __trailing = [c for c in trailing if c in df.columns]
+        __leading = [c for c in leading if c in columns]
+        __trailing = [c for c in trailing if c in columns]
 
-        columns = __leading + fields + __trailing
+        return __leading + fields + __trailing
 
+    def reorder_dataframe(self, df: DataFrame) -> DataFrame:
+        columns = self.sort_columns(df.columns)
+        columns = [f"`{c}`" for c in columns]
         return df.select(columns)
 
     @abstractmethod
-    def optimize_table(self):
-        raise NotImplementedError()
+    def optimize_table(self): ...
 
     @abstractmethod
-    def update_schema(self, src: Union[DataFrame, Table, str], **kwargs):
-        raise NotImplementedError()
+    def update_schema(self, src: AllowedSources, **kwargs): ...
 
     @abstractmethod
-    def get_differences_with_deltatable(self, src: Union[DataFrame, Table, str], **kwargs):
-        raise NotImplementedError()
+    def get_differences_with_deltatable(self, src: AllowedSources, **kwargs): ...
 
     @abstractmethod
-    def overwrite_schema(self, src: Union[DataFrame, Table, str]):
-        raise NotImplementedError()
+    def overwrite_schema(self, src: AllowedSources): ...
 
     def __str__(self):
         return f"{self.table.qualified_name}"

fabricks/cdc/base/generator.py

@@ -4,11 +4,12 @@ from typing import Any, List, Optional, Sequence, Union, cast
 
 from py4j.protocol import Py4JJavaError
 from pyspark.sql import DataFrame
-from pyspark.sql.connect.dataframe import DataFrame as CDataFrame
 
+from fabricks.cdc.base._types import AllowedSources
 from fabricks.cdc.base.configurator import Configurator
 from fabricks.context.log import DEFAULT_LOGGER
 from fabricks.metastore.table import SchemaDiff, Table
+from fabricks.utils._types import DataFrameLike
 from fabricks.utils.sqlglot import fix as fix_sql
 
 
@@ -18,13 +19,17 @@ class Generator(Configurator):
 
     def create_table(
         self,
-        src: Union[DataFrame, Table, str],
+        src: AllowedSources,
         partitioning: Optional[bool] = False,
         partition_by: Optional[Union[List[str], str]] = None,
         identity: Optional[bool] = False,
         liquid_clustering: Optional[bool] = False,
         cluster_by: Optional[Union[List[str], str]] = None,
         properties: Optional[dict[str, str]] = None,
+        masks: Optional[dict[str, str]] = None,
+        primary_key: Optional[dict[str, Any]] = None,
+        foreign_keys: Optional[dict[str, Any]] = None,
+        comments: Optional[dict[str, str]] = None,
         **kwargs,
     ):
         kwargs["mode"] = "complete"
@@ -37,7 +42,7 @@ class Generator(Configurator):
         if partitioning is True:
             assert partition_by, "partitioning column(s) not found"
 
-        df = self.reorder_columns(df)
+        df = self.reorder_dataframe(df)
 
         identity = False if identity is None else identity
         liquid_clustering = False if liquid_clustering is None else liquid_clustering
@@ -50,16 +55,20 @@ class Generator(Configurator):
             liquid_clustering=liquid_clustering,
             cluster_by=cluster_by,
             properties=properties,
+            masks=masks,
+            primary_key=primary_key,
+            foreign_keys=foreign_keys,
+            comments=comments,
         )
 
     def create_or_replace_view(self, src: Union[Table, str], schema_evolution: bool = True, **kwargs):
-        assert not isinstance(src, (DataFrame, CDataFrame)), "dataframe not allowed"
+        assert not isinstance(src, DataFrameLike), "dataframe not allowed"
 
         assert kwargs["mode"] == "complete", f"{kwargs['mode']} not allowed"
         sql = self.get_query(src, **kwargs)
 
         df = self.spark.sql(sql)
-        df = self.reorder_columns(df)
+        df = self.reorder_dataframe(df)
         columns = [f"`{c}`" for c in df.columns]
 
         sql = f"""
@@ -74,12 +83,12 @@ class Generator(Configurator):
         from __view
         """
         sql = fix_sql(sql)
-        DEFAULT_LOGGER.debug("create or replace view", extra={"job": self, "sql": sql})
+        DEFAULT_LOGGER.debug("create or replace view", extra={"label": self, "sql": sql})
 
         try:
             self.spark.sql(sql)
-        except Py4JJavaError:
-            DEFAULT_LOGGER.exception("could not execute sql query", extra={"job": self, "sql": sql})
+        except Py4JJavaError as e:
+            DEFAULT_LOGGER.exception("fail to execute sql query", extra={"label": self, "sql": sql}, exc_info=e)
 
     def optimize_table(self):
         columns = None
@@ -91,35 +100,34 @@ class Generator(Configurator):
 
         self.table.optimize(columns=columns)
 
-    def get_differences_with_deltatable(self, src: Union[DataFrame, Table, str], **kwargs) -> Optional[DataFrame]:
+    def get_differences_with_deltatable(self, src: AllowedSources, **kwargs) -> DataFrame:
+        from pyspark.sql.types import StringType, StructField, StructType
+
+        schema = StructType(
+            [
+                StructField("column", StringType(), False),
+                StructField("data_type", StringType(), True),
+                StructField("new_column", StringType(), True),
+                StructField("new_data_type", StringType(), True),
+                StructField("status", StringType(), True),
+            ]
+        )
+
         if self.is_view:
-            return None
+            return self.spark.createDataFrame([], schema=schema)
 
         else:
-            from pyspark.sql.types import StringType, StructField, StructType
-
             kwargs["mode"] = "complete"
             if "slice" in kwargs:
                 del kwargs["slice"]
 
             df = self.get_data(src, **kwargs)
-            df = self.reorder_columns(df)
+            df = self.reorder_dataframe(df)
+
             diffs = self.table.get_schema_differences(df)
-            df_diff = self.spark.createDataFrame(
-                [cast(Any, d.model_dump()) for d in diffs],
-                schema=StructType(
-                    [
-                        StructField("column", StringType(), False),
-                        StructField("data_type", StringType(), True),
-                        StructField("new_column", StringType(), True),
-                        StructField("new_data_type", StringType(), True),
-                        StructField("status", StringType(), True),
-                    ]
-                ),
-            )
-            return df_diff
-
-    def get_schema_differences(self, src: Union[DataFrame, Table, str], **kwargs) -> Optional[Sequence[SchemaDiff]]:
+            return self.spark.createDataFrame([cast(Any, d.model_dump()) for d in diffs], schema=schema)
+
+    def get_schema_differences(self, src: AllowedSources, **kwargs) -> Optional[Sequence[SchemaDiff]]:
         if self.is_view:
             return None
 
@@ -129,10 +137,11 @@ class Generator(Configurator):
                 del kwargs["slice"]
 
             df = self.get_data(src, **kwargs)
-            df = self.reorder_columns(df)
+            df = self.reorder_dataframe(df)
+
             return self.table.get_schema_differences(df)
 
-    def schema_drifted(self, src: Union[DataFrame, Table, str], **kwargs) -> Optional[bool]:
+    def schema_drifted(self, src: AllowedSources, **kwargs) -> Optional[bool]:
         d = self.get_schema_differences(src, **kwargs)
         if d is None:
             return None
@@ -140,13 +149,13 @@ class Generator(Configurator):
 
     def _update_schema(
         self,
-        src: Union[DataFrame, Table, str],
+        src: AllowedSources,
         overwrite: bool = False,
         widen_types: bool = False,
         **kwargs,
     ):
         if self.is_view:
-            assert not isinstance(src, (DataFrame, CDataFrame)), "dataframe not allowed"
+            assert not isinstance(src, DataFrameLike), "dataframe not allowed"
             self.create_or_replace_view(src=src)
 
         else:
@@ -155,14 +164,14 @@ class Generator(Configurator):
                 del kwargs["slice"]
 
             df = self.get_data(src, **kwargs)
-            df = self.reorder_columns(df)
+            df = self.reorder_dataframe(df)
             if overwrite:
                 self.table.overwrite_schema(df)
             else:
                 self.table.update_schema(df, widen_types=widen_types)
 
-    def update_schema(self, src: Union[DataFrame, Table, str], **kwargs):
+    def update_schema(self, src: AllowedSources, **kwargs):
         self._update_schema(src=src, **kwargs)
 
-    def overwrite_schema(self, src: Union[DataFrame, Table, str], **kwargs):
+    def overwrite_schema(self, src: AllowedSources, **kwargs):
         self._update_schema(src=src, overwrite=True, **kwargs)