fabricks 2024.7.1.5__py3-none-any.whl

fabricks/utils/azure_queue.py

@@ -0,0 +1,63 @@
+import json
+from typing import Optional, Union
+
+from azure.core.exceptions import ResourceExistsError
+from azure.storage.queue import QueueClient
+
+
+class AzureQueue(QueueClient):
+    def __init__(
+        self,
+        name: str,
+        storage_account: Optional[str] = None,
+        access_key: Optional[str] = None,
+        connection_string: Optional[str] = None,
+    ):
+        self.name = name
+        if connection_string is None:
+            assert storage_account
+            assert access_key
+            self.storage_account = storage_account
+            self.access_key = access_key
+            connection_string = f"DefaultEndpointsProtocol=https;AccountName={self.storage_account};AccountKey={self.access_key};EndpointSuffix=core.windows.net"
+
+        assert connection_string
+        self.connection_string = connection_string
+
+    @property
+    def queue_client(self) -> QueueClient:
+        return QueueClient.from_connection_string(self.connection_string, queue_name=self.name)
+
+    def create_if_not_exists(self):
+        try:
+            self.queue_client.create_queue()
+        except ResourceExistsError:
+            pass
+
+    @property
+    def sentinel(self):
+        return "SENTINEL"
+
+    def clear(self):
+        self.queue_client.clear_messages()
+
+    def send(self, message: Union[str, dict]):
+        if isinstance(message, dict):
+            message = json.dumps(message)
+        # print("sending ->", message)
+        self.queue_client.send_message(message)
+
+    def send_sentinel(self):
+        # print("sentinel", self.sentinel)
+        self.send(self.sentinel)
+
+    def receive(self):
+        msg = self.queue_client.receive_message()
+        if msg:
+            self.queue_client.delete_message(msg)
+            # print("receiving ->", msg.content)
+            return msg.content
+        return None
+
+    def delete(self):
+        self.queue_client.delete_queue()

fabricks/utils/azure_table.py

@@ -0,0 +1,99 @@
+import time
+from typing import List, Optional, Union
+
+from azure.data.tables import TableClient, TableServiceClient
+from pyspark.sql import DataFrame
+
+
+class AzureTable(TableClient):
+    def __init__(
+        self,
+        name: str,
+        storage_account: Optional[str] = None,
+        access_key: Optional[str] = None,
+        connection_string: Optional[str] = None,
+    ):
+        self.name = name
+        if connection_string is None:
+            assert storage_account
+            assert access_key
+            self.storage_account = storage_account
+            self.access_key = access_key
+            connection_string = f"DefaultEndpointsProtocol=https;AccountName={self.storage_account};AccountKey={self.access_key};EndpointSuffix=core.windows.net"
+
+        assert connection_string
+        self.connection_string = connection_string
+
+    @property
+    def table_service_client(self) -> TableServiceClient:
+        return TableServiceClient.from_connection_string(self.connection_string, table_name=self.name)
+
+    @property
+    def table(self) -> TableClient:
+        return self.create_if_not_exists()
+
+    def create_if_not_exists(self) -> TableClient:
+        return self.table_service_client.create_table_if_not_exists(table_name=self.name)
+
+    def drop(self):
+        self.table_service_client.delete_table(self.name)
+
+    def query(self, query: str) -> List:
+        return list(self.table.query_entities(query))
+
+    def list_all(self) -> List:
+        return self.query("")
+
+    def submit(self, operations: List, retry: Optional[bool] = True):
+        try:
+            partitions = set()
+            for d in operations:
+                partitions.add(d[1]["PartitionKey"])
+
+            for p in partitions:
+                _operations = [d for d in operations if d[1].get("PartitionKey") == p]
+                t = 50
+                if len(_operations) < t:
+                    self.table.submit_transaction(_operations)
+                else:
+                    transactions = [_operations[i : i + t] for i in range(0, len(_operations), t)]
+                    for transaction in transactions:
+                        self.table.submit_transaction(transaction)
+        except Exception as e:
+            if retry:
+                time.sleep(10)
+                self.submit(operations, retry=False)
+            else:
+                raise e
+
+    def delete(self, data: Union[List, DataFrame, dict]):
+        if isinstance(data, DataFrame):
+            data = [row.asDict() for row in data.collect()]
+        elif not isinstance(data, List):
+            data = [data]
+
+        operations = [("delete", d) for d in data]
+        self.submit(operations)
+
+    def upsert(self, data: Union[List, DataFrame, dict]):
+        if isinstance(data, DataFrame):
+            data = [row.asDict() for row in data.collect()]
+        elif not isinstance(data, List):
+            data = [data]
+
+        operations = [("upsert", d) for d in data]
+        self.submit(operations)
+
+    def truncate_partition(self, partition: str):
+        data = self.query(f"PartitionKey eq '{partition}'")
+        self.delete(data)
+
+    def truncate_all_partitions(self):
+        for p in self.list_all_partitions():
+            self.truncate_partition(p)
+
+    def list_all_partitions(self) -> List:
+        partitions = set()
+        for d in self.list_all():
+            partitions.add(d["PartitionKey"])
+        return sorted(list(partitions))

fabricks/utils/console.py

@@ -0,0 +1,51 @@
+from typing import Optional
+
+
+class formatter:
+    END = "\33[0m"
+    BOLD = "\33[1m"
+    ITALIC = "\33[3m"
+    URL = "\33[4m"
+    BLINK = "\33[5m"
+    SELECTED = "\33[7m"
+
+    BLINK2 = "\33[6m"
+
+
+class colors:
+    BLACK = "\33[30m"
+    RED = "\33[31m"
+    GREEN = "\33[32m"
+    YELLOW = "\33[33m"
+    BLUE = "\33[34m"
+    VIOLET = "\33[35m"
+    BEIGE = "\33[36m"
+    WHITE = "\33[37m"
+    GREY = "\33[90m"
+    ORANGE = "\33[33m"
+
+    RED2 = "\33[91m"
+    GREEN2 = "\33[92m"
+    YELLOW2 = "\33[93m"
+    BLUE2 = "\33[94m"
+    VIOLET2 = "\33[95m"
+    BEIGE2 = "\33[96m"
+    WHITE2 = "\33[97m"
+
+    RED3 = "\33[1;31m"
+
+
+def progress_bar(progress: int = 0, width: int = 40, msg: Optional[str] = None):
+    if not isinstance(progress, int):
+        progress = int(progress)
+
+    left = width * progress // 100
+    right = width - left
+
+    tags = "#" * left
+    spaces = " " * right
+    pct = f" {progress}%"
+    if msg:
+        pct = f"{pct} ({msg})"
+
+    print("\r[", tags, spaces, "]", pct, sep="", end="", flush=True)

fabricks/utils/container.py

@@ -0,0 +1,57 @@
+from databricks.sdk.runtime import dbutils
+
+from fabricks.utils.secret import AccessKey, ApplicationRegistration, Secret
+
+
+def create_container(storage_account: str, container: str, secret: Secret):
+    from azure.core.exceptions import ResourceExistsError
+    from azure.identity import ClientSecretCredential
+    from azure.storage.blob import BlobServiceClient
+
+    assert isinstance(secret, ApplicationRegistration)
+
+    cred = ClientSecretCredential(
+        tenant_id=secret.directory_id,
+        client_id=secret.application_id,
+        client_secret=secret.secret,
+    )
+
+    try:
+        blob_service_client = BlobServiceClient(f"https://{storage_account}.blob.core.windows.net", credential=cred)
+        blob_service_client.create_container(container)
+
+    except ResourceExistsError:
+        pass
+
+
+def mount_container(storage_account: str, container: str, secret: Secret):
+    try:
+        dbutils.fs.unmount(f"/mnt/{container}")  # type: ignore
+
+    except Exception:
+        pass
+
+    if isinstance(secret, AccessKey):
+        dbutils.fs.mount(  # type: ignore
+            source=f"wasbs://{container}@{storage_account}.blob.core.windows.net",
+            mount_point=f"/mnt/{container}",
+            extra_configs={
+                f"fs.azure.account.key.{storage_account}.blob.core.windows.net": f"{secret.key}",
+            },
+        )
+
+    elif isinstance(secret, ApplicationRegistration):
+        dbutils.fs.mount(  # type: ignore
+            source=f"abfss://{container}@{storage_account}.blob.core.windows.net/",
+            mount_point=f"/mnt/{container}",
+            extra_configs={
+                "fs.azure.account.auth.type": "OAuth",
+                "fs.azure.account.oauth.provider.type": "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredentialsTokenProvider",
+                "fs.azure.account.oauth2.client.id": secret.application_id,
+                "fs.azure.account.oauth2.client.secret": secret.secret,
+                "fs.azure.account.oauth2.client.endpoint": f"https://login.microsoftonline.com/{secret.directory_id}/oauth2/token",
+            },
+        )
+
+    else:
+        raise ValueError("secret is not valid")

fabricks/utils/fdict.py

@@ -0,0 +1,28 @@
+from typing import Any, List, Optional
+
+
+class FDict:
+    def __init__(self, options: Any):
+        self.options = options
+
+    def get(self, key: str) -> Optional[Any]:
+        return self.options.get(key)
+
+    def get_list(self, key: str) -> List[Any]:
+        values = self.options.get(key, [])
+        if values is not None:
+            if not isinstance(values, List):
+                values = [values]
+        return values
+
+    def get_boolean(self, key: str, if_none: Optional[bool] = None) -> Optional[bool]:
+        o = self.options.get(key)
+        if isinstance(o, bool):
+            return o
+        elif o is not None:
+            return o.lower() == "true"
+        else:
+            return if_none
+
+    def get_dict(self, key) -> dict:
+        return self.options.get(key, {})

fabricks/utils/helpers.py

@@ -0,0 +1,89 @@
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from functools import reduce
+from typing import Any, Callable, Iterable, List, Optional, Union
+
+from databricks.sdk.runtime import dbutils, spark
+from pyspark.sql import DataFrame
+from typing_extensions import deprecated
+
+from fabricks.utils.path import Path
+
+
+def concat_ws(fields: Union[str, List[str]], alias: Optional[str] = None) -> str:
+    if isinstance(fields, str):
+        fields = [fields]
+
+    if alias:
+        coalesce = [f"coalesce(cast({alias}.{f} as string), '-1')" for f in fields]
+    else:
+        coalesce = [f"coalesce(cast({f} as string), '-1')" for f in fields]
+
+    return "concat_ws('*', " + ",".join(coalesce) + ")"
+
+
+def concat_dfs(dfs: Iterable[DataFrame]) -> DataFrame:
+    return reduce(lambda x, y: x.unionByName(y, allowMissingColumns=True), dfs)
+
+
+@deprecated("use run_threads instead")
+def run_threads(func: Callable, iter: Union[List, DataFrame, range, set], workers: int = 8) -> List[Any]:
+    return run_in_parallel(func, iter, workers)
+
+
+def run_in_parallel(func: Callable, iterable: Union[List, DataFrame, range, set], workers: int = 8) -> List[Any]:
+    """
+    Runs the given function in parallel on the elements of the iterable using multiple threads.
+
+    Args:
+        func (Callable): The function to be executed in parallel.
+        iterable (Union[List, DataFrame, range, set]): The iterable containing the elements on which the function will be executed.
+        workers (int, optional): The number of worker threads to use. Defaults to 8.
+
+    Returns:
+        List[Any]: A list containing the results of the function calls.
+
+    """
+    out = []
+
+    with ThreadPoolExecutor(max_workers=workers) as executor:
+        iterable = iterable.collect() if isinstance(iterable, DataFrame) else iterable
+        futures = {executor.submit(func, i): i for i in iterable}
+        for future in as_completed(futures):
+            try:
+                r = future.result()
+                if r:
+                    out.append(r)
+            except Exception:
+                pass
+
+    return out
+
+
+def run_notebook(path: Path, timeout: Optional[int] = None, **kwargs):
+    """
+    Runs a notebook located at the given path.
+
+    Args:
+        path (Path): The path to the notebook file.
+        timeout (Optional[int]): The maximum execution time for the notebook in seconds. Defaults to None.
+        **kwargs: Additional keyword arguments to be passed to the notebook.
+
+    Returns:
+        None
+    """
+    if timeout is None:
+        timeout = 3600
+
+    dbutils.notebook.run(path.get_notebook_path(), timeout, {**kwargs})
+
+
+def xxhash64(s: Any):
+    df = spark.sql(f"select xxhash64(cast('{s}' as string)) as xxhash64")
+    return df.collect()[0][0]
+
+
+def md5(s: Any):
+    from hashlib import md5
+
+    md5 = md5(str(s).encode())
+    return md5.hexdigest()

fabricks/utils/log.py

@@ -0,0 +1,153 @@
+import hashlib
+import json
+import logging
+from datetime import datetime
+from typing import Tuple
+
+from fabricks.utils.azure_table import AzureTable
+
+GREY = "\33[90m"
+BLACK = "\33[30m"
+YELLOW = "\33[93m"
+RED = "\33[31m"
+RESET = "\33[0m"
+COLORS = {"DEBUG": GREY, "INFO": BLACK, "WARNING": YELLOW, "ERROR": RED, "CRITICAL": RED}
+
+
+class CustomFormatter(logging.Formatter):
+    def format(self, record):
+        message = super().format(record)  # noqa: F841
+        out = f"{COLORS[record.levelname]}"
+
+        if hasattr(record, "created"):
+            d = datetime.fromtimestamp(record.created).strftime("%d/%m/%y %H:%M:%S")
+            out += f"[{d}]"
+
+        if hasattr(record, "job"):
+            j = f" - {record.__dict__.get('job')}"
+            out += str(j)
+
+        elif hasattr(record, "step"):
+            s = f" - {record.__dict__.get('step')}"
+            out += str(s)
+
+        if hasattr(record, "message"):
+            m = record.__dict__.get("message", "")
+            if hasattr(record, "job"):
+                m = f" => {m}"
+            elif hasattr(record, "step"):
+                m = f" => {m}"
+            else:
+                m = " " + m
+            out += m
+
+        if hasattr(record, "exc_info"):
+            exc_info = record.__dict__.get("exc_info", None)
+            if exc_info is not None:
+                e = f" !{exc_info[0].__name__.lower()}!"
+                out += e
+
+        if hasattr(record, "sql"):
+            s = f"\n---\n%sql\n{record.__dict__.get('sql')}\n---"
+            out += s
+        if hasattr(record, "content"):
+            s = f"\n---\n{record.__dict__.get('content')}\n---"
+            out += s
+
+        out += f"{RESET}"
+        return out
+
+
+class AzureTableHandler(logging.Handler):
+    def __init__(self, table: AzureTable):
+        super().__init__()
+        self.buffer = []
+        self.table = table
+
+    def emit(self, record):
+        if hasattr(record, "target"):
+            target = record.__dict__.get("target")
+
+            r = {
+                "Created": str(
+                    datetime.fromtimestamp(record.created).strftime("%d/%m/%y %H:%M:%S")
+                ),  # timestamp not present when querying Azure Table
+                "Level": record.levelname,
+                "Message": record.message,
+            }
+            if hasattr(record, "job"):
+                j = str(record.__dict__.get("job", ""))
+                r["Job"] = j
+                r["JobId"] = hashlib.md5(j.encode()).hexdigest()
+            if hasattr(record, "table"):
+                t = str(record.__dict__.get("table", ""))
+                r["Job"] = t
+                r["JobId"] = hashlib.md5(t.encode()).hexdigest()
+
+            if hasattr(record, "step"):
+                r["Step"] = record.__dict__.get("step", "")
+
+            if hasattr(record, "schedule_id"):
+                r["ScheduleId"] = record.__dict__.get("schedule_id", "")
+
+            if hasattr(record, "schedule"):
+                r["Schedule"] = record.__dict__.get("schedule", "")
+
+            if hasattr(record, "notebook_id"):
+                r["NotebookId"] = record.__dict__.get("notebook_id", "")
+
+            if hasattr(record, "exc_info"):
+                e = record.__dict__.get("exc_info", None)
+                if e is not None:
+                    d = {
+                        "type": str(e[0].__name__)[:1000],
+                        "message": str(e[1])[:1000],
+                        "traceback": str(logging.Formatter.formatException(self, e))[:1000],  # type: ignore
+                    }
+                    r["Exception"] = json.dumps(d)
+
+            if hasattr(record, "content"):
+                r["Content"] = json.dumps(record.__dict__.get("content", ""))[:1000]
+            if hasattr(record, "sql"):
+                r["Sql"] = record.__dict__.get("sql", "")[:1000]
+
+            r["PartitionKey"] = record.__dict__.get("partition_key", "default")
+            if hasattr(record, "row_key"):
+                r["RowKey"] = record.__dict__.get("row_key", "")
+            else:
+                r["RowKey"] = hashlib.md5(json.dumps(r, sort_keys=True).encode()).hexdigest()
+
+            if target == "table":
+                self.table.upsert(r)
+            else:
+                self.buffer.append(r)
+
+        else:
+            pass
+
+    def flush(self):
+        self.table.upsert(self.buffer)
+        self.buffer.clear()
+
+    def clear_buffer(self):
+        self.buffer = []
+
+
+def get_logger(name: str, level: int, table: AzureTable) -> Tuple[logging.Logger, AzureTableHandler]:
+    logger = logging.getLogger(name)
+    logger.setLevel(level)
+
+    # Console handler
+    console_handler = logging.StreamHandler()
+    console_handler.setLevel(level)
+    console_format = CustomFormatter()
+    console_handler.setFormatter(console_format)
+
+    # Azure Table handler
+    azure_table_handler = AzureTableHandler(table=table)
+    azure_table_handler.setLevel(level)
+
+    logger.addHandler(console_handler)
+    logger.addHandler(azure_table_handler)
+
+    return logger, azure_table_handler