fabricks 3.0.11__py3-none-any.whl

fabricks/utils/helpers.py

@@ -0,0 +1,228 @@
+import logging
+from functools import reduce
+from queue import Queue
+from typing import Any, Callable, Iterable, List, Literal, Optional, Union
+
+from pyspark.sql import DataFrame
+from typing_extensions import deprecated
+
+from fabricks.utils._types import DataFrameLike
+from fabricks.utils.path import Path
+from fabricks.utils.spark import spark
+
+
+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]) -> Optional[DataFrame]:
+    dfs = [df for df in dfs if df is not None]
+    if len(dfs) == 0:
+        return None
+    return reduce(lambda x, y: x.unionByName(y, allowMissingColumns=True), dfs)
+
+
+@deprecated("use run_in_parallel 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 _process_queue_item(func: Callable, task_queue: Queue, result_queue: Queue, stop_signal: Any):
+    """Worker function that processes items from a queue."""
+    while True:
+        try:
+            item = task_queue.get(timeout=1)
+
+            if item is stop_signal:
+                task_queue.put(stop_signal)  # Put it back for other workers
+                break
+
+            result = func(item)
+            result_queue.put(result)
+        except Exception:
+            continue
+
+
+def _run_in_parallel_legacy(
+    func: Callable,
+    iterable: Union[List, DataFrame, range, set],
+    workers: int = 8,
+    progress_bar: Optional[bool] = False,
+    position: Optional[int] = None,
+) -> List[Any]:
+    from concurrent.futures import ThreadPoolExecutor
+
+    iterable = iterable.collect() if isinstance(iterable, DataFrameLike) else iterable  # type: ignore
+
+    with ThreadPoolExecutor(max_workers=workers) as executor:
+        if progress_bar:
+            from tqdm import tqdm
+
+            results = list(tqdm(executor.map(func, iterable), total=len(iterable), position=position))
+        else:
+            results = list(executor.map(func, iterable))
+
+    return results
+
+
+def run_in_parallel(
+    func: Callable,
+    iterable: Union[List, DataFrame, range, set],
+    workers: int = 8,
+    progress_bar: Optional[bool] = False,
+    position: Optional[int] = None,
+    loglevel: int = logging.CRITICAL,
+    logger: Optional[logging.Logger] = None,
+    run_as: Optional[Literal["ThreadPool", "ProcessPool", "Pool", "Queue", "Legacy"]] = "Legacy",
+) -> List[Any]:
+    """
+    Runs the given function in parallel on the elements of the iterable using multiple threads or processes.
+
+    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/processes to use. Defaults to 8.
+        progress_bar (Optional[bool], optional): Whether to display a progress bar. Defaults to False.
+        position (Optional[int], optional): Position for the progress bar. Defaults to None.
+        loglevel (int, optional): Log level to set during execution. Defaults to logging.CRITICAL.
+        logger (Optional[logging.Logger], optional): Logger instance to use. Defaults to None.
+        run_as (Optional[Literal["ThreadPool", "ProcessPool", "Pool", "Queue"]], optional): Type of run as to use.
+
+    Returns:
+        List[Any]: A list containing the results of the function calls.
+
+    """
+    if logger is None:
+        logger = logging.getLogger()
+
+    current_loglevel = logger.getEffectiveLevel()
+    logger.setLevel(loglevel)
+
+    if run_as == "Legacy":
+        results = _run_in_parallel_legacy(
+            func=func,
+            iterable=iterable,
+            workers=workers,
+            progress_bar=progress_bar,
+            position=position,
+        )
+
+    else:
+        iterables = iterable.collect() if isinstance(iterable, DataFrameLike) else iterable  # type: ignore
+        results = []
+
+        if run_as == "Queue":
+            import threading
+
+            task_queue = Queue()
+            result_queue = Queue()
+            stop_signal = object()
+
+            for item in iterables:
+                task_queue.put(item)
+
+            task_queue.put(stop_signal)
+
+            threads = []
+            for _ in range(workers):
+                t = threading.Thread(target=_process_queue_item, args=(func, task_queue, result_queue, stop_signal))
+                t.start()
+
+                threads.append(t)
+
+            if progress_bar:
+                from tqdm import tqdm
+
+                with tqdm(total=len(iterables), position=position) as t:
+                    for _ in range(len(iterables)):
+                        result = result_queue.get()
+                        results.append(result)
+
+                        t.update()
+                        t.refresh()
+
+            else:
+                for _ in range(len(iterables)):
+                    results.append(result_queue.get())
+
+            for t in threads:
+                t.join()
+
+        elif run_as == "Pool":
+            from multiprocessing import Pool
+
+            with Pool(processes=workers) as p:
+                if progress_bar:
+                    from tqdm import tqdm
+
+                    with tqdm(total=len(iterables), position=position) as t:
+                        for result in p.map(func, iterables):
+                            results.append(result)
+
+                            t.update()
+                            t.refresh()
+
+                else:
+                    results = list(p.map(func, iterables))
+
+        else:
+            from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
+
+            Executor = ProcessPoolExecutor if run_as == "ProcessPool" else ThreadPoolExecutor
+            with Executor(max_workers=workers) as exe:
+                if progress_bar:
+                    from tqdm import tqdm
+
+                    with tqdm(total=len(iterables), position=position) as t:
+                        for result in exe.map(func, iterables):
+                            results.append(result)
+
+                            t.update()
+                            t.refresh()
+
+                else:
+                    results = list(exe.map(func, iterables))
+
+    logger.setLevel(current_loglevel)
+
+    return results
+
+
+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
+    """
+    from databricks.sdk.runtime import dbutils
+
+    if timeout is None:
+        timeout = 3600
+
+    dbutils.notebook.run(path.get_notebook_path(), timeout, {**kwargs})  # type: ignore
+
+
+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,236 @@
+import hashlib
+import json
+import logging
+import sys
+from datetime import datetime
+from datetime import timezone as tz
+from typing import Optional, Tuple
+from zoneinfo import ZoneInfo
+
+from pyspark.sql import DataFrame
+
+from fabricks.utils.azure_table import AzureTable
+
+
+class LogFormatter(logging.Formatter):
+    def __init__(self, debugmode: Optional[bool] = False, timezone: Optional[str] = None):
+        super().__init__(fmt="%(levelname)s%(prefix)s%(message)s [%(timestamp)s]%(extra)s")
+
+        self.debugmode = False if debugmode is None else debugmode
+        self.timezone = ZoneInfo(timezone) if timezone else tz.utc
+
+    COLORS = {
+        logging.DEBUG: "\033[36m",
+        logging.INFO: "\033[32m",
+        logging.WARNING: "\033[33m",
+        logging.ERROR: "\033[31m",
+        logging.CRITICAL: "\033[41;31m",
+    }
+
+    RESET = "\033[0m"
+    BRIGHT = "\033[1m"
+
+    PADDINGS = {
+        "DEBUG": "   ",
+        "INFO": "    ",
+        "WARNING": " ",
+        "ERROR": "   ",
+        "CRITICAL": "",
+    }
+
+    def formatTime(self, record) -> str:
+        ct = datetime.fromtimestamp(record.created, tz=tz.utc).astimezone(self.timezone)
+        s = ct.strftime("%d/%m/%y %H:%M:%S")
+        return f"{self.COLORS[logging.DEBUG]}{s}{self.RESET}"
+
+    def format(self, record):
+        levelname = record.levelname
+        padding = self.PADDINGS[levelname]
+        levelname_formatted = f"{self.COLORS[record.levelno]}{levelname}:{padding}{self.RESET}"
+
+        prefix = ""
+
+        if hasattr(record, "label"):
+            prefix = f"{record.__dict__.get('label')} - "
+        elif hasattr(record, "job"):
+            prefix = f"{record.__dict__.get('job')} - "  # keep for backward compatibility
+        elif hasattr(record, "step"):
+            prefix = f"{self.BRIGHT}{record.__dict__.get('step')}{self.RESET} - "
+
+        extra = ""
+        if hasattr(record, "exc_info") and record.exc_info:
+            exc_info = record.__dict__.get("exc_info", None)
+            extra += f" [{self.COLORS[logging.ERROR]}{exc_info[0].__name__}{self.RESET}]"
+
+        if self.debugmode:
+            if hasattr(record, "sql"):
+                extra += f"\n---\n%sql\n{record.__dict__.get('sql')}\n---"
+
+            if hasattr(record, "content"):
+                extra += f"\n---\n{record.__dict__.get('content')}\n---"
+
+            if hasattr(record, "context"):
+                extra += f"\n---\n{json.dumps(record.__dict__.get('context'), indent=2, default=str)}\n---"
+
+            if hasattr(record, "df"):
+                df = record.__dict__.get("df")
+                if isinstance(df, DataFrame):
+                    extra += f"\n---\n%df\n{df.toPandas().to_string(index=True)}\n---"
+
+        record.levelname = levelname_formatted
+        record.prefix = prefix
+        record.timestamp = self.formatTime(record)
+        record.extra = extra
+
+        return super().format(record)
+
+
+class AzureTableLogHandler(logging.Handler):
+    def __init__(self, table: AzureTable, debugmode: Optional[bool] = False, timezone: Optional[str] = None):
+        super().__init__()
+
+        self.buffer = []
+        self.table = table
+
+        self.debugmode = False if debugmode is None else debugmode
+        self.timezone = ZoneInfo(timezone) if timezone else tz.utc
+
+    def formatTime(self, record) -> str:
+        ct = datetime.fromtimestamp(record.created, tz=tz.utc).astimezone(self.timezone)
+        s = ct.strftime("%d/%m/%y %H:%M:%S")
+        return s
+
+    def emit(self, record):
+        if hasattr(record, "target"):
+            target = record.__dict__.get("target")
+
+            level = record.levelname
+            if "debug" in level.lower():
+                level = "DEBUG"
+            elif "info" in level.lower():
+                level = "INFO"
+            elif "warning" in level.lower():
+                level = "WARNING"
+            elif "error" in level.lower():
+                level = "ERROR"
+            elif "critical" in level.lower():
+                level = "CRITICAL"
+            else:
+                level = "INFO"
+
+            r = {
+                "Created": self.formatTime(record),  # timestamp not present when querying Azure Table
+                "Level": level,
+                "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 self.debugmode:
+                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 = []
+
+
+class CustomConsoleHandler(logging.StreamHandler):
+    def __init__(self, stream=None, debugmode: Optional[bool] = False):
+        super().__init__(stream or sys.stderr)
+
+        self.debugmode = False if debugmode is None else debugmode
+
+    def emit(self, record):
+        if hasattr(record, "sql"):
+            if self.debugmode:
+                super().emit(record)
+        else:
+            super().emit(record)
+
+
+def get_logger(
+    name: str,
+    level: int,
+    table: Optional[AzureTable] = None,
+    debugmode: Optional[bool] = False,
+    timezone: Optional[str] = None,
+) -> Tuple[logging.Logger, Optional[AzureTableLogHandler]]:
+    logger = logging.getLogger(name)
+    if logger.hasHandlers():
+        logger.handlers.clear()
+
+    root = logging.getLogger()
+    if root.hasHandlers():
+        root.handlers.clear()
+
+    logger.setLevel(level)
+    logger.propagate = False
+
+    # Console handler
+    console_handler = CustomConsoleHandler(debugmode=debugmode)
+    console_handler.setLevel(level)
+    console_format = LogFormatter(debugmode=debugmode, timezone=timezone)
+    console_handler.setFormatter(console_format)
+
+    if table is not None:
+        # Azure Table handler
+        azure_table_handler = AzureTableLogHandler(table=table, debugmode=debugmode, timezone=timezone)
+        azure_table_handler.setLevel(level)
+    else:
+        azure_table_handler = None
+
+    logger.addHandler(console_handler)
+    if azure_table_handler is not None:
+        logger.addHandler(azure_table_handler)
+
+    return logger, azure_table_handler

fabricks/utils/mermaid.py

@@ -0,0 +1,32 @@
+from pyspark.sql import DataFrame
+
+
+def get_mermaid_diagram(df: DataFrame) -> str:
+    dependencies = df.select("parent_id", "parent", "job_id", "job").collect()
+
+    out = "flowchart TD\n"
+
+    unique_nodes = set()
+
+    for row in dependencies:
+        parent_id = str(row["parent_id"])
+        parent_name = str(row["parent"])
+        child_id = str(row["job_id"])
+        child_name = str(row["job"])
+
+        if parent_id != "0" and parent_id is not None:
+            if parent_id not in unique_nodes:
+                out += f"    {parent_id}[{parent_name}]\n"
+                unique_nodes.add(parent_id)
+
+            if child_id not in unique_nodes:
+                out += f"    {child_id}[{child_name}]\n"
+                unique_nodes.add(child_id)
+
+            out += f"    {parent_id} --> {child_id}\n"
+        else:
+            if child_id not in unique_nodes:
+                out += f"    {child_id}[{child_name}]\n"
+                unique_nodes.add(child_id)
+
+    return out