fabricks 3.0.11__py3-none-any.whl

fabricks/utils/read/read_excel.py

@@ -0,0 +1,5 @@
+from pyspark.sql import DataFrame
+
+
+def read_excel(path: str) -> DataFrame:
+    raise NotImplementedError()

fabricks/utils/read/read_yaml.py

@@ -0,0 +1,33 @@
+from typing import Iterable, Optional, cast
+
+import yaml
+
+from fabricks.utils.path import Path
+
+
+def read_yaml(
+    path: Path,
+    root: Optional[str] = None,
+    preferred_file_name: Optional[str] = None,
+) -> Iterable[dict]:
+    found = False
+
+    for file in path.walk():
+        if not file.endswith(".yml"):
+            continue
+
+        if preferred_file_name is not None and preferred_file_name not in file:
+            continue
+
+        found = True
+
+        with open(file, "r", encoding="utf-8") as f:
+            data = yaml.safe_load(f)
+            for job_config in data:
+                if root:
+                    yield cast(dict, job_config[root])
+                else:
+                    yield cast(dict, job_config)
+
+    if preferred_file_name is not None and not found:
+        yield from read_yaml(path=path, root=root, preferred_file_name=None)

fabricks/utils/schema/__init__.py

@@ -0,0 +1,7 @@
+from fabricks.utils.schema.get_json_schema_for_type import get_json_schema_for_type
+from fabricks.utils.schema.get_schema_for_type import get_schema_for_type
+
+__all__ = [
+    "get_json_schema_for_type",
+    "get_schema_for_type",
+]

fabricks/utils/schema/get_json_schema_for_type.py

@@ -0,0 +1,161 @@
+import dataclasses
+import datetime
+import logging
+import sys
+import types
+from typing import Any, ForwardRef, Literal, Type, Union, get_type_hints
+from uuid import UUID
+
+LOGGER = logging.getLogger(__name__)
+
+
+def get_json_schema_for_type(proptype: Type):
+    def_list: dict[str, dict] = {}
+    schema = _get_json_schema_for_type(proptype, def_list, is_root=True)
+    schema["$defs"] = def_list
+    schema["$schema"] = "https://json-schema.org/draft/2020-12/schema"
+    return schema
+
+
+def _get_json_schema_for_type(proptype: Type, def_list: dict[str, dict], is_root: bool, is_nullable=False) -> dict:
+    def _fixref(input: dict) -> dict:
+        if "type" in input:
+            if "$ref" in input["type"]:
+                return input["type"]
+        return input
+
+    def _may_null(input: dict, is_nullable: bool) -> dict:
+        if is_nullable:
+            return {"oneOf": [{"type": "null"}, input]}
+        return input
+
+    if hasattr(proptype, "__origin__") and proptype.__origin__ == Literal:
+        return {"enum": proptype.__args__}
+
+    if hasattr(proptype, "__origin__") and proptype.__origin__ == tuple:  # noqa E721
+        return {
+            "type": "array",
+            "minItems": len(proptype.__args__),
+            "maxItems": len(proptype.__args__),
+            "additionalItems": False,
+            "prefixItems": [_get_json_schema_for_type(t, def_list, is_root=False) for t in proptype.__args__],
+        }
+
+    if (sys.version_info >= (3, 10) and isinstance(proptype, types.UnionType)) or (
+        hasattr(proptype, "__origin__") and proptype.__origin__ == Union
+    ):
+        if len(proptype.__args__) == 2 and proptype.__args__[0] == type(None):  # noqa E721
+            t = _get_json_schema_for_type(proptype.__args__[1], def_list, is_root=False, is_nullable=True)
+            return t
+
+        if len(proptype.__args__) == 2 and proptype.__args__[1] == type(None):  # noqa E721
+            t = _get_json_schema_for_type(proptype.__args__[0], def_list, is_root=False, is_nullable=True)
+            return t
+
+        one_of_types = [
+            _get_json_schema_for_type(f, def_list, is_root=False, is_nullable=False) for f in proptype.__args__
+        ]
+
+        return {"oneOf": one_of_types}
+
+    if proptype == type(None):  # noqa E721
+        return {"type": "null"}
+
+    if proptype == str:  # noqa E721
+        return {"type": "string"} if not is_nullable else {"type": ["string", "null"]}
+
+    if proptype == Any:
+        return {}
+
+    if proptype == UUID:
+        return {
+            "type": "string" if not is_nullable else ["string", "null"],
+            "format": "uuid",
+        }
+
+    if proptype == int:  # noqa E721
+        return {"type": "integer" if not is_nullable else ["integer", "null"]}
+
+    if proptype == float:  # noqa E721
+        return {"type": "number" if not is_nullable else ["number", "null"]}
+
+    if proptype == bool:  # noqa E721
+        return {"type": "boolean" if not is_nullable else ["boolean", "null"]}
+
+    if hasattr(proptype, "__origin__") and proptype.__origin__ == list:  # noqa E721
+        return {
+            "type": "array",
+            "items": _get_json_schema_for_type(proptype.__args__[0], def_list, is_root=False),
+        }
+
+    if hasattr(proptype, "__bases__") and len(proptype.__bases__) == 1 and proptype.__bases__[0] == dict:  # noqa E721
+        typehints = get_type_hints(proptype)
+        props = {k: _get_json_schema_for_type(v, def_list, is_root=False) for (k, v) in typehints.items()}
+
+        if hasattr(proptype, "__name__") and not is_root:
+            def_list[proptype.__name__] = {"type": "object", "properties": props}
+            return _may_null({"$ref": "#/$defs/" + proptype.__name__}, is_nullable)
+        else:
+            return _may_null({"type": "object", "properties": props}, is_nullable)
+
+    if dataclasses.is_dataclass(proptype):
+        required = [
+            f.name
+            for f in dataclasses.fields(proptype)
+            if f.default == dataclasses.MISSING and f.default_factory == dataclasses.MISSING and f.init
+        ]
+        definition = {
+            "type": "object",
+            "required": required,
+            "additionalProperties": False,
+            "properties": {
+                f.name: _get_json_schema_for_type(f.type, def_list, is_root=False)  # type: ignore
+                for f in dataclasses.fields(proptype)
+            },
+        }
+
+        if is_root:
+            return definition
+        else:
+            def_list[proptype.__name__] = definition
+
+            return _may_null({"$ref": "#/$defs/" + proptype.__name__}, is_nullable)
+
+    if hasattr(proptype, "__origin__") and proptype.__origin__ == dict and len(proptype.__args__) == 2:  # noqa E721
+        keytype = proptype.__args__[0]
+        if keytype != str and keytype != UUID:  # noqa E721
+            raise NotImplementedError()
+        valuetype = proptype.__args__[1]
+        return _may_null(
+            {
+                "type": "object",
+                "additionalProperties": _fixref(
+                    {"type": _get_json_schema_for_type(valuetype, def_list, is_root=False)}
+                ),
+            },
+            is_nullable,
+        )
+
+    if isinstance(proptype, ForwardRef):
+        arg = proptype.__forward_arg__
+        return _may_null({"$ref": "#/$defs/" + arg}, is_nullable)
+
+    if proptype == datetime.datetime:
+        return {
+            "type": "string" if not is_nullable else ["string", "null"],
+            "format": "date-time",
+        }
+
+    if proptype == datetime.time:
+        return {
+            "type": "string" if not is_nullable else ["string", "null"],
+            "format": "time",
+        }
+
+    if proptype == datetime.date:
+        return {
+            "type": "string" if not is_nullable else ["string", "null"],
+            "format": "date",
+        }
+
+    return {}

fabricks/utils/schema/get_schema_for_type.py

@@ -0,0 +1,99 @@
+import dataclasses
+from typing import List, Literal, Type, Union, cast, get_type_hints, overload
+
+from pyspark.sql.types import (
+    ArrayType,
+    BooleanType,
+    DataType,
+    DoubleType,
+    LongType,
+    MapType,
+    NullType,
+    StringType,
+    StructField,
+    StructType,
+)
+
+
+@overload
+def get_schema_for_type(proptype: Union[int, str, float, bool]) -> DataType: ...
+
+
+@overload
+def get_schema_for_type(proptype: Type) -> StructType: ...
+
+
+def _merge_struct_types(types: List[DataType]):
+    not_none_types = [t for t in types if type(t) != type(NullType())]  # noqa: E721
+
+    assert len([f for f in not_none_types if not isinstance(f, StructType)]) == 0
+    all_fields: List[StructField] = []
+
+    for subtype in not_none_types:
+        fields = cast(StructType, subtype).fields
+        for field in fields:
+            existing_field = next((f for f in all_fields if f.name == field.name), None)
+            if existing_field is not None and (
+                type(existing_field.dataType) != type(field.dataType)  # noqa: E721
+                or isinstance(existing_field.dataType, StructType)
+            ):
+                new_type = _merge_struct_types([existing_field.dataType, field.dataType])
+                all_fields.append(StructField(name=field.name, dataType=new_type))
+                all_fields.remove(existing_field)
+            else:
+                assert existing_field is None or type(existing_field.dataType) == type(field.dataType)  # noqa: E721
+                if existing_field is None:
+                    all_fields.append(field)
+
+    return StructType(fields=all_fields)
+
+
+def get_schema_for_type(proptype: Type) -> DataType:  # type: ignore
+    if hasattr(proptype, "__origin__") and proptype.__origin__ == Literal:
+        return get_schema_for_type(type(proptype.__args__[0]))  # For literal types we assume first type is correct
+
+    if hasattr(proptype, "__origin__") and proptype.__origin__ == Union:
+        if len(proptype.__args__) == 2 and proptype.__args__[0] == type(None):  # noqa E721
+            return get_schema_for_type(proptype.__args__[1])
+        if len(proptype.__args__) == 2 and proptype.__args__[1] == type(None):  # noqa E721
+            return get_schema_for_type(proptype.__args__[0])
+
+        return _merge_struct_types([get_schema_for_type(f) for f in proptype.__args__])
+
+    if proptype == type(None):  # noqa E721
+        return NullType()
+
+    if proptype == str:  # noqa E721
+        return StringType()
+
+    if proptype == int:  # noqa E721
+        return LongType()
+
+    if proptype == float:  # noqa E721
+        return DoubleType()
+
+    if proptype == bool:  # noqa E721
+        return BooleanType()
+
+    if hasattr(proptype, "__origin__") and proptype.__origin__ == list:  # noqa E721
+        return ArrayType(get_schema_for_type(proptype.__args__[0]))
+
+    if proptype == dict[str, str]:
+        return MapType(StringType(), StringType())
+
+    if hasattr(proptype, "__bases__") and len(proptype.__bases__) == 1 and proptype.__bases__[0] == dict:  # noqa E721
+        types = get_type_hints(proptype)
+        fields = [StructField(k, get_schema_for_type(v)) for k, v in types.items()]
+        return StructType(fields=fields)
+
+    if dataclasses.is_dataclass(proptype):
+        fields = [StructField(f.name, get_schema_for_type(f.type)) for f in dataclasses.fields(proptype)]
+        return StructType(fields=fields)
+
+    if hasattr(proptype, "__origin__") and proptype.__origin__ == dict:  # noqa E721
+        if len(proptype.__args__) == 2:
+            value_type = proptype.__args__[1]
+            value_schema = get_schema_for_type(value_type)
+            return MapType(StringType(), value_schema)
+
+    raise NotImplementedError()

fabricks/utils/spark.py

@@ -0,0 +1,76 @@
+import os
+from typing import Final, Optional
+
+from databricks.sdk.dbutils import RemoteDbUtils
+from pyspark.sql import DataFrame, SparkSession
+
+DATABRICKS_LOCALMODE: Final[bool] = os.getenv("DATABRICKS_LOCALMODE", "false").lower() in ("true", "1", "yes")
+
+
+def get_spark() -> SparkSession:
+    if DATABRICKS_LOCALMODE:
+        from databricks.connect.session import DatabricksSession
+        from databricks.sdk.core import Config
+
+        profile = os.getenv("DATABRICKS_PROFILE", "DEFAULT")
+
+        cluster_id = os.getenv("DATABRICKS_CLUSTER_ID")
+        assert cluster_id, "DATABRICKS_CLUSTER_ID environment variable is not set"
+
+        c = Config(profile=profile, cluster_id=cluster_id)
+
+        spark = DatabricksSession.builder.sdkConfig(c).getOrCreate()
+
+    else:
+        pass
+
+        spark = SparkSession.builder.getOrCreate()  # type: ignore
+
+    assert spark is not None
+    return spark  # type: ignore
+
+
+def display(df: DataFrame, limit: Optional[int] = None) -> None:
+    """
+    Display a Spark DataFrame in Databricks notebook or local environment.
+    If running in local mode, it converts the DataFrame to a Pandas DataFrame for display.
+    """
+    if DATABRICKS_LOCALMODE:
+        from IPython.display import display
+
+        if limit is not None:
+            df = df.limit(limit)
+
+        display(df.toPandas())
+
+    else:
+        from databricks.sdk.runtime import display
+
+        if limit is not None:
+            df = df.limit(limit)
+
+        display(df)
+
+
+def get_dbutils(spark: Optional[SparkSession] = None) -> Optional[RemoteDbUtils]:
+    try:
+        if DATABRICKS_LOCALMODE:
+            from databricks.sdk import WorkspaceClient
+
+            w = WorkspaceClient()
+            dbutils = w.dbutils
+
+        else:
+            from pyspark.dbutils import DBUtils
+
+            dbutils = DBUtils(spark)
+
+        assert dbutils is not None
+        return dbutils  # type: ignore
+
+    except Exception:
+        return None
+
+
+spark = get_spark()
+dbutils = get_dbutils(spark=spark)

fabricks/utils/sqlglot.py

@@ -0,0 +1,56 @@
+from typing import List, Optional
+
+from sqlglot import exp, parse_one, transpile
+from sqlglot.dialects.databricks import Databricks
+
+
+class Fabricks(Databricks):
+    class Generator(Databricks.Generator):
+        EXPRESSIONS_WITHOUT_NESTED_CTES = {
+            exp.Insert,
+            exp.Union,
+        }
+
+
+def fix(sql: str, keep_comments: bool = True):
+    sql = transpile(
+        sql,
+        "fabricks",
+        identify=True,
+        pretty=True,
+        normalize=False,
+        normalize_functions="lower",
+        leading_comma=True,
+        max_text_width=119,
+        comments=keep_comments,
+    )[0]
+    return sql
+
+
+def is_global_temp_view(sql: str):
+    tables = parse_one(sql, dialect="fabricks").find_all(exp.Table)
+    for t in tables:
+        return "global_temp" in str(t)
+
+
+def get_global_temp_view(sql: str) -> Optional[str]:
+    tables = parse_one(sql, dialect="fabricks").find_all(exp.Table)
+    for t in tables:
+        if "global_temp" in str(t):
+            return str(t)
+
+
+def parse(sql: str) -> exp.Expression:
+    return parse_one(sql, dialect="fabricks")
+
+
+def get_tables(sql: str, allowed_databases: Optional[List[str]] = None) -> List[str]:
+    tables = set()
+    for table in parse(sql).find_all(exp.Table):
+        if len(table.db) > 0:  # exclude CTEs
+            if allowed_databases:
+                if table.db not in allowed_databases:
+                    continue
+            tables.add(f"{table.db}.{table.name}")
+    tables = list(tables)
+    return tables

fabricks/utils/write/__init__.py

@@ -0,0 +1,8 @@
+from fabricks.utils.write.delta import append_delta, overwrite_delta
+from fabricks.utils.write.stream import write_stream
+
+__all__ = [
+    "append_delta",
+    "overwrite_delta",
+    "write_stream",
+]

fabricks/utils/write/delta.py

@@ -0,0 +1,46 @@
+from typing import List, Optional, Union, get_args
+
+from pyspark.sql import DataFrame
+
+from fabricks.utils.path import Path
+from fabricks.utils.read._types import AllowedIOModes
+
+
+def write_delta(
+    df: DataFrame,
+    path: Path,
+    mode: AllowedIOModes,
+    options: Optional[dict[str, str]] = None,
+    partition_by: Union[Optional[List[str]], str] = None,
+):
+    assert mode in list(get_args(AllowedIOModes))
+
+    if isinstance(partition_by, str):
+        partition_by = [partition_by]
+
+    writer = df.write.format("delta").mode(mode).option("mergeSchema", "True").option("overwriteSchema", "True")
+    if partition_by:
+        writer = writer.partitionBy(*partition_by)
+    if options:
+        for key, value in options.items():
+            writer = writer.option(key, value)
+
+    writer.save(path.string)
+
+
+def append_delta(
+    df: DataFrame,
+    path: Path,
+    options: Optional[dict[str, str]] = None,
+    partition_by: Union[Optional[List[str]], str] = None,
+):
+    write_delta(df, path, "append", options=options)
+
+
+def overwrite_delta(
+    df: DataFrame,
+    path: Path,
+    options: Optional[dict[str, str]] = None,
+    partition_by: Union[Optional[List[str]], str] = None,
+):
+    write_delta(df, path, "overwrite", options=options)

fabricks/utils/write/stream.py

@@ -0,0 +1,27 @@
+from typing import Callable, Optional
+
+from pyspark.sql import DataFrame
+from pyspark.sql.streaming.query import StreamingQuery
+
+from fabricks.utils.path import Path
+
+
+def write_stream(
+    df: DataFrame,
+    checkpoints_path: Path,
+    func: Callable,
+    timeout: Optional[int] = 18000,
+) -> StreamingQuery:
+    if timeout is None:
+        timeout = 18000
+
+    assert timeout is not None
+
+    query = (
+        df.writeStream.foreachBatch(func)
+        .option("checkpointLocation", checkpoints_path.string)
+        .trigger(once=True)
+        .start()
+    )
+    query.awaitTermination(timeout=timeout)
+    return query

fabricks-3.0.11.dist-info/METADATA

@@ -0,0 +1,23 @@
+Metadata-Version: 2.4
+Name: fabricks
+Version: 3.0.11
+Author-email: BMS DWH Team <bi_support@bmsuisse.ch>
+Requires-Python: <4,>=3.9
+Requires-Dist: azure-data-tables<13,>=12.5.0
+Requires-Dist: azure-identity>=1.10.0
+Requires-Dist: azure-storage-blob>=12.14.1
+Requires-Dist: azure-storage-queue<13,>=12.10.0
+Requires-Dist: databricks-cli>=0.18.0
+Requires-Dist: databricks-sdk>=0.20.0
+Requires-Dist: importlib-metadata>=8.6.1
+Requires-Dist: jinja2>=2.11.3
+Requires-Dist: mermaid-magic>=0.1.4
+Requires-Dist: pydantic-settings
+Requires-Dist: pydantic-yaml>=1.4.0
+Requires-Dist: pydantic>=1.10.21
+Requires-Dist: python-dotenv>=1.0.1
+Requires-Dist: pyyaml>=6.0.0
+Requires-Dist: sqlglot>=22.1.1
+Requires-Dist: tenacity>=9.1.2
+Requires-Dist: tomli>=2.2.1; python_version < '3.11'
+Requires-Dist: tqdm>=4.67.1