fabricks 3.0.11__py3-none-any.whl

fabricks/core/jobs/base/_types.py

@@ -0,0 +1,284 @@
+from dataclasses import dataclass
+from typing import List, Literal, Optional, TypedDict, Union
+
+from pydantic import BaseModel, ConfigDict, model_validator
+from pyspark.sql.types import StringType, StructField, StructType
+
+from fabricks.cdc.base._types import AllowedChangeDataCaptures
+from fabricks.context import BRONZE, GOLD, SILVER
+from fabricks.core.jobs.get_job_id import get_dependency_id, get_job_id
+from fabricks.core.parsers import ParserOptions
+from fabricks.utils.fdict import FDict
+from fabricks.utils.path import Path
+
+TBronze = Literal["bronze"]
+TSilver = Literal["silver"]
+TGold = Literal["gold"]
+TStep = Literal[TBronze, TSilver, TGold]
+
+Bronzes: List[TBronze] = [b.get("name") for b in BRONZE]
+Silvers: List[TSilver] = [s.get("name") for s in SILVER]
+Golds: List[TGold] = [g.get("name") for g in GOLD]
+Steps: List[TStep] = Bronzes + Silvers + Golds
+
+AllowedModesBronze = Literal["memory", "append", "register"]
+AllowedModesSilver = Literal["memory", "append", "latest", "update", "combine"]
+AllowedModesGold = Literal["memory", "append", "complete", "update", "invoke"]
+AllowedModes = Literal[AllowedModesBronze, AllowedModesSilver, AllowedModesGold]
+
+AllowedFileFormats = Literal["json_array", "json", "jsonl", "csv", "parquet", "delta"]
+AllowedOperations = Literal["upsert", "reload", "delete"]
+AllowedTypes = Literal["manual", "default"]
+AllowedOrigins = Literal["parser", "job"]
+
+AllowedConstraintOptions = Literal["not enforced", "deferrable", "initially deferred", "norely", "rely"]
+AllowedForeignKeyOptions = Literal["match full", "on update no action", "on delete no action"]
+
+
+class SparkOptions(TypedDict):
+    sql: Optional[dict[str, str]]
+    conf: Optional[dict[str, str]]
+
+
+class ForeignKeyOptions(TypedDict):
+    foreign_key: Optional[AllowedForeignKeyOptions]
+    constraint: Optional[AllowedConstraintOptions]
+
+
+class PrimaryKeyOptions(TypedDict):
+    constraint: Optional[AllowedConstraintOptions]
+
+
+class ForeignKey(TypedDict):
+    keys: List[str]
+    reference: str
+    options: Optional[ForeignKeyOptions]
+
+
+class PrimaryKey(TypedDict):
+    keys: List[str]
+    options: Optional[PrimaryKeyOptions]
+
+
+class TableOptions(TypedDict):
+    identity: Optional[bool]
+    liquid_clustering: Optional[bool]
+    partition_by: Optional[List[str]]
+    zorder_by: Optional[List[str]]
+    cluster_by: Optional[List[str]]
+    powerbi: Optional[bool]
+    maximum_compatibility: Optional[bool]
+    bloomfilter_by: Optional[List[str]]
+    constraints: Optional[dict[str, str]]
+    properties: Optional[dict[str, str]]
+    comment: Optional[str]
+    calculated_columns: Optional[dict[str, str]]
+    masks: Optional[dict[str, str]]
+    comments: Optional[dict[str, str]]
+    retention_days: Optional[int]
+    primary_key: Optional[dict[str, PrimaryKey]]
+    foreign_keys: Optional[dict[str, ForeignKey]]
+
+
+class _InvokeOptions(TypedDict):
+    notebook: str
+    timeout: int
+    arguments: Optional[dict[str, str]]
+
+
+class InvokerOptions(TypedDict):
+    pre_run: Optional[List[_InvokeOptions]]
+    run: Optional[List[_InvokeOptions]]
+    post_run: Optional[List[_InvokeOptions]]
+
+
+class ExtenderOptions(TypedDict):
+    extender: str
+    arguments: Optional[dict[str, str]]
+
+
+class CheckOptions(TypedDict):
+    skip: Optional[bool]
+    pre_run: Optional[bool]
+    post_run: Optional[bool]
+    min_rows: Optional[int]
+    max_rows: Optional[int]
+    count_must_equal: Optional[str]
+
+
+class BronzeOptions(TypedDict):
+    type: Optional[AllowedTypes]
+    mode: AllowedModesBronze
+    uri: str
+    parser: str
+    source: str
+    keys: Optional[List[str]]
+    # default
+    parents: Optional[List[str]]
+    filter_where: Optional[str]
+    optimize: Optional[bool]
+    compute_statistics: Optional[bool]
+    vacuum: Optional[bool]
+    no_drop: Optional[bool]
+    # extra
+    encrypted_columns: Optional[List[str]]
+    calculated_columns: Optional[dict[str, str]]
+    operation: Optional[AllowedOperations]
+    timeout: Optional[int]
+
+
+class SilverOptions(TypedDict):
+    type: Optional[AllowedTypes]
+    mode: AllowedModesSilver
+    change_data_capture: AllowedChangeDataCaptures
+    # default
+    parents: Optional[List[str]]
+    filter_where: Optional[str]
+    optimize: Optional[bool]
+    compute_statistics: Optional[bool]
+    vacuum: Optional[bool]
+    no_drop: Optional[bool]
+    # extra
+    deduplicate: Optional[bool]
+    stream: Optional[bool]
+    # else
+    order_duplicate_by: Optional[dict[str, str]]
+    timeout: Optional[int]
+
+
+class GoldOptions(TypedDict):
+    type: Optional[AllowedTypes]
+    mode: AllowedModesGold
+    change_data_capture: AllowedChangeDataCaptures
+    update_where: Optional[str]
+    # default
+    parents: Optional[List[str]]
+    optimize: Optional[bool]
+    compute_statistics: Optional[bool]
+    vacuum: Optional[bool]
+    no_drop: Optional[bool]
+    # extra
+    deduplicate: Optional[bool]  # remove duplicates on the keys and on the hash
+    rectify_as_upserts: Optional[bool]  # convert reloads into upserts and deletes
+    correct_valid_from: Optional[bool]  # update valid_from to '1900-01-01' for the first timestamp
+    persist_last_timestamp: Optional[bool]  # persist the last timestamp to be used as a watermark for the next run
+    # delete_missing: Optional[bool]  # delete missing records on update (to be implemented)
+    # else
+    table: Optional[str]
+    notebook: Optional[bool]
+    requirements: Optional[bool]
+    timeout: Optional[int]
+    metadata: Optional[bool]
+
+
+StepOptions = Union[BronzeOptions, SilverOptions, GoldOptions]
+
+
+@dataclass
+class BaseJobConf:
+    job_id: str
+    topic: str
+    item: str
+
+
+@dataclass
+class JobConfBronze(BaseJobConf):
+    step: TBronze
+    options: BronzeOptions
+    table_options: Optional[TableOptions] = None
+    parser_options: Optional[ParserOptions] = None
+    check_options: Optional[CheckOptions] = None
+    spark_options: Optional[SparkOptions] = None
+    invoker_options: Optional[InvokerOptions] = None
+    extender_options: Optional[List[ExtenderOptions]] = None
+    tags: Optional[List[str]] = None
+    comment: Optional[str] = None
+
+
+@dataclass
+class JobConfSilver(BaseJobConf):
+    step: TSilver
+    options: SilverOptions
+    table_options: Optional[TableOptions] = None
+    check_options: Optional[CheckOptions] = None
+    spark_options: Optional[SparkOptions] = None
+    invoker_options: Optional[InvokerOptions] = None
+    extender_options: Optional[List[ExtenderOptions]] = None
+    tags: Optional[List[str]] = None
+    comment: Optional[str] = None
+
+
+@dataclass
+class JobConfGold(BaseJobConf):
+    step: TGold
+    options: Optional[GoldOptions]
+    table_options: Optional[TableOptions] = None
+    check_options: Optional[CheckOptions] = None
+    spark_options: Optional[SparkOptions] = None
+    invoker_options: Optional[InvokerOptions] = None
+    extender_options: Optional[List[ExtenderOptions]] = None
+    tags: Optional[List[str]] = None
+    comment: Optional[str] = None
+
+
+JobConf = Union[JobConfBronze, JobConfSilver, JobConfGold]
+
+
+@dataclass
+class Paths:
+    storage: Path
+    tmp: Path
+    checkpoints: Path
+    commits: Path
+    schema: Path
+    runtime: Path
+
+
+@dataclass
+class Options:
+    job: FDict
+    check: FDict
+    table: FDict
+    spark: FDict
+    invokers: FDict
+    extenders: List
+
+
+class JobDependency(BaseModel):
+    model_config = ConfigDict(extra="forbid", frozen=True)
+    origin: AllowedOrigins
+    job_id: str
+    parent: str
+    parent_id: str
+    dependency_id: str
+
+    def __str__(self) -> str:
+        return f"{self.job_id} -> {self.parent}"
+
+    @model_validator(mode="after")
+    def check_no_circular_dependency(self):
+        if self.job_id == self.parent_id:
+            raise ValueError("Circular dependency detected")
+        return self
+
+    @staticmethod
+    def from_parts(job_id: str, parent: str, origin: AllowedOrigins):
+        parent = parent.removesuffix("__current")
+        return JobDependency(
+            job_id=job_id,
+            origin=origin,
+            parent=parent,
+            parent_id=get_job_id(job=parent),
+            dependency_id=get_dependency_id(parent=parent, job_id=job_id),
+        )
+
+
+SchemaDependencies = StructType(
+    [
+        StructField("dependency_id", StringType(), True),
+        StructField("origin", StringType(), True),
+        StructField("job_id", StringType(), True),
+        StructField("parent_id", StringType(), True),
+        StructField("parent", StringType(), True),
+    ]
+)

fabricks/core/jobs/base/checker.py

@@ -0,0 +1,139 @@
+from typing import Literal
+
+from fabricks.context.log import DEFAULT_LOGGER
+from fabricks.core.jobs.base.exception import (
+    PostRunCheckException,
+    PostRunCheckWarning,
+    PreRunCheckException,
+    PreRunCheckWarning,
+    SkipRunCheckWarning,
+)
+from fabricks.core.jobs.base.generator import Generator
+
+
+class Checker(Generator):
+    def check_pre_run(self):
+        self._check("pre_run")
+
+    def check_post_run(self):
+        self._check("post_run")
+
+    def _check(self, position: Literal["pre_run", "post_run"]):
+        if self.options.check.get(position):
+            DEFAULT_LOGGER.debug(f"check {position}", extra={"label": self})
+
+            p = self.paths.runtime.append(f".{position}.sql")
+            assert p.exists(), f"{position} check not found ({p})"
+
+            df = self.spark.sql(p.get_sql())
+            fail_df = df.where("__action == 'fail'")
+            warning_df = df.where("__action == 'warning'")
+
+            if not fail_df.isEmpty():
+                for row in fail_df.collect():
+                    DEFAULT_LOGGER.warning(
+                        f"check {position} failed due to {row['__message']}",
+                        extra={"label": self},
+                    )
+
+                if position == "pre_run":
+                    raise PreRunCheckException(row["__message"], dataframe=df)
+                elif position == "post_run":
+                    raise PostRunCheckException(row["__message"], dataframe=df)
+
+            elif not warning_df.isEmpty():
+                for row in warning_df.collect():
+                    DEFAULT_LOGGER.warning(
+                        f"check {position} failed due to {row['__message']}",
+                        extra={"label": self},
+                    )
+
+                if position == "pre_run":
+                    raise PreRunCheckWarning(row["__message"], dataframe=df)
+                elif position == "post_run":
+                    raise PostRunCheckWarning(row["__message"], dataframe=df)
+
+    def check_post_run_extra(self):
+        min_rows = self.options.check.get("min_rows")
+        max_rows = self.options.check.get("max_rows")
+        count_must_equal = self.options.check.get("count_must_equal")
+
+        if min_rows or max_rows or count_must_equal:
+            df = self.spark.sql(f"select count(*) from {self}")
+            rows = df.collect()[0][0]
+            if min_rows:
+                DEFAULT_LOGGER.debug("check min rows", extra={"label": self})
+                if rows < min_rows:
+                    raise PostRunCheckException(f"min rows check failed ({rows} < {min_rows})", dataframe=df)
+
+            if max_rows:
+                DEFAULT_LOGGER.debug("check max rows", extra={"label": self})
+                if rows > max_rows:
+                    raise PostRunCheckException(f"max rows check failed ({rows} > {max_rows})", dataframe=df)
+
+            if count_must_equal:
+                DEFAULT_LOGGER.debug("check count must equal", extra={"label": self})
+                equals_rows = self.spark.read.table(count_must_equal).count()
+                if rows != equals_rows:
+                    raise PostRunCheckException(
+                        f"count must equal check failed ({count_must_equal} - {rows} != {equals_rows})",
+                        dataframe=df,
+                    )
+
+    def _check_duplicate_in_column(self, column: str):
+        if column in self.table.columns:
+            DEFAULT_LOGGER.debug(f"check duplicate in {column}", extra={"label": self})
+
+            cols = [column]
+
+            if "__source" in self.table.columns:
+                cols.append("__source")
+
+            if self.change_data_capture == "scd2":
+                cols.append("__valid_to")
+
+            elif self.change_data_capture == "nocdc":
+                if "__valid_to" in self.table.columns:
+                    cols.append("__valid_to")
+                elif self.mode == "append" and "__timestamp" in self.table.columns:
+                    cols.append("__timestamp")
+
+            cols = ", ".join(cols)
+            df = self.spark.sql(f"select {cols} from {self} group by all having count(*) > 1 limit 5")
+
+            if not df.isEmpty():
+                duplicates = ",".join([str(row[column]) for row in df.collect()])
+                raise PostRunCheckException(
+                    f"duplicate {column} check failed ({duplicates})",
+                    dataframe=df,
+                )
+
+        else:
+            DEFAULT_LOGGER.debug(f"could not find {column}", extra={"label": self})
+
+    def check_duplicate_key(self):
+        self._check_duplicate_in_column("__key")
+
+    def check_duplicate_hash(self):
+        self._check_duplicate_in_column("__hash")
+
+    def check_duplicate_identity(self):
+        self._check_duplicate_in_column("__identity")
+
+    def check_skip_run(self):
+        if self.options.check.get("skip"):
+            DEFAULT_LOGGER.debug("check if run should be skipped", extra={"label": self})
+
+            p = self.paths.runtime.append(".skip.sql")
+            assert p.exists(), "skip check not found"
+
+            df = self.spark.sql(p.get_sql())
+            skip_df = df.where("__skip")
+            if not skip_df.isEmpty():
+                for row in skip_df.collect():
+                    DEFAULT_LOGGER.warning(
+                        f"skip run due to {row['__message']}",
+                        extra={"label": self},
+                    )
+
+                raise SkipRunCheckWarning(row["__message"], dataframe=df)