fabricks 2024.7.1.5__py3-none-any.whl

fabricks/utils/read/read_yaml.py

@@ -0,0 +1,43 @@
+from typing import Optional
+
+import yaml
+from databricks.sdk.runtime import spark
+from pyspark.sql import DataFrame
+from pyspark.sql.types import StructType
+
+from fabricks.utils.helpers import concat_dfs
+from fabricks.utils.path import Path
+
+
+def read_yaml(
+    path: Path,
+    root: Optional[str] = None,
+    schema: Optional[StructType] = None,
+    file_name: Optional[str] = None,
+) -> Optional[DataFrame]:
+    files = [f for f in path.walk() if f.endswith(".yml")]
+    if file_name:
+        files = [f for f in files if file_name in f]
+
+    dfs = [spark.createDataFrame([], schema=schema)] if schema else []
+
+    for file in files:
+        with open(file) as f:
+            data = yaml.safe_load(f)
+
+        if schema:
+            dt = [d[root] for d in data] if root else data
+            df = spark.createDataFrame(dt, schema=schema)
+        else:
+            json = spark.sparkContext.parallelize(data)
+            df = spark.read.json(json)
+            if root:
+                df = df.select(f"{root}.*")
+
+        dfs.append(df)
+
+    if dfs:
+        df = concat_dfs(dfs)
+        return df
+
+    return spark.createDataFrame([], schema=schema) if schema else None

fabricks/utils/read/types.py

@@ -0,0 +1,3 @@
+from typing import Literal
+
+IOModes = Literal["overwrite", "append"]

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):
+            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):
+            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):
+        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,93 @@
+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):
+            return get_schema_for_type(proptype.__args__[1])
+        if len(proptype.__args__) == 2 and proptype.__args__[1] == type(None):
+            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):
+        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)
+
+    raise NotImplementedError()

fabricks/utils/secret.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+
+from databricks.sdk.runtime import dbutils, spark
+
+
+@dataclass
+class Secret:
+    pass
+
+
+@dataclass
+class ApplicationRegistration(Secret):
+    secret: str
+    application_id: str
+    directory_id: str
+
+
+@dataclass
+class AccessKey(Secret):
+    key: str
+
+
+def _get_secret_from_secret_scope(secret_scope: str, name: str) -> str:
+    scopes = [s.name for s in dbutils.secrets.listScopes()]
+    assert secret_scope in scopes, "scope {secret_scope} not found"
+    return dbutils.secrets.get(scope=secret_scope, key=name)
+
+
+def get_secret_from_secret_scope(secret_scope: str, name: str) -> Secret:
+    secret = _get_secret_from_secret_scope(secret_scope=secret_scope, name=name)
+    if name.endswith("application-registration"):
+        s = json.loads(secret)
+        assert s.get("secret"), f"no secret found in {name}"
+        assert s.get("application_id"), f"no application_id found in {name}"
+        assert s.get("directory_id"), f"no directory_id found in {name}"
+        return ApplicationRegistration(
+            secret=s.get("secret"),
+            application_id=s.get("application_id"),
+            directory_id=s.get("directory_id"),
+        )
+    elif name.endswith("access-key"):
+        return AccessKey(key=secret)
+    else:
+        raise ValueError(f"{name} is not valid")
+
+
+def _add_secret_to_spark(key: str, value: str):
+    spark.conf.set(key, value)
+    # needed for check (invalid configuration value detected for fs.azure.account.key)
+    spark._jsc.hadoopConfiguration().set(key, value)  # type: ignore
+
+
+def add_secret_to_spark(secret: Secret, uri: str):
+    if isinstance(secret, ApplicationRegistration):
+        _add_secret_to_spark(f"fs.azure.account.auth.type.{uri}", "OAuth")
+        _add_secret_to_spark(
+            f"fs.azure.account.oauth.provider.type.{uri}",
+            "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider",
+        )
+        _add_secret_to_spark(
+            f"fs.azure.account.oauth2.client.id.{uri}",
+            secret.application_id,
+        )
+        _add_secret_to_spark(
+            f"fs.azure.account.oauth2.client.secret.{uri}",
+            secret.secret,
+        )
+        _add_secret_to_spark(
+            f"fs.azure.account.oauth2.client.endpoint.{uri}",
+            f"https://login.microsoftonline.com/{secret.directory_id}/oauth2/token",
+        )
+    elif isinstance(secret, AccessKey):
+        _add_secret_to_spark(f"fs.azure.account.key.{uri}", secret.key)
+    else:
+        raise ValueError("secret is not valid")

fabricks/utils/sqlglot.py

@@ -0,0 +1,48 @@
+from typing import 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):
+    """
+    Fixes the given SQL query by parsing it using the 'fabricks' dialect,
+    transpiling it, and returning the fixed SQL query.
+
+    Args:
+        sql (str): The SQL query to be fixed.
+
+    Returns:
+        str: The fixed SQL query.
+    """
+    sql = parse_one(sql, dialect="fabricks").sql()
+    sql = transpile(
+        sql,
+        identify=True,
+        pretty=True,
+        normalize=False,
+        normalize_functions="lower",
+        write="fabricks",
+    )[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)

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 IOModes
+
+
+def write_delta(
+    df: DataFrame,
+    path: Path,
+    mode: IOModes,
+    options: Optional[dict[str, str]] = None,
+    partition_by: Union[Optional[List[str]], str] = None,
+):
+    assert mode in list(get_args(IOModes))
+
+    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