fabricks 3.0.11__py3-none-any.whl

fabricks/metastore/utils.py

@@ -0,0 +1,51 @@
+from pyspark.errors.exceptions.base import AnalysisException
+from pyspark.sql import DataFrame
+
+from fabricks.context import SPARK
+
+
+def get_tables(schema: str) -> DataFrame:
+    table_df = SPARK.sql(f"show tables in {schema}")
+    view_df = SPARK.sql(f"show views in {schema}")
+
+    try:
+        df = SPARK.sql(
+            """
+            select 
+              database,
+              concat_ws('.', database, tableName) as table,
+              md5(table) as job_id
+            from 
+              {tables} 
+              left anti join {views} on tableName = viewName
+            """,
+            tables=table_df,
+            views=view_df,
+        )
+        return df
+
+    except AnalysisException:
+        return SPARK.sql("select null::string as database, null::string as table")
+
+
+def get_views(schema: str) -> DataFrame:
+    view_df = SPARK.sql(f"show views in {schema}")
+
+    try:
+        df = SPARK.sql(
+            """
+            select 
+              namespace as database,
+              concat_ws('.', namespace, viewName) as view,
+              md5(view) as job_id
+            from 
+                {views}
+            where
+              not isTemporary
+            """,
+            views=view_df,
+        )
+        return df
+
+    except AnalysisException:
+        return SPARK.sql("select null::string as database, null::string as view")

fabricks/metastore/view.py

@@ -0,0 +1,53 @@
+from typing import Any, Optional, Union
+from uuid import uuid4
+
+import pandas as pd
+from pyspark.sql import DataFrame, SparkSession
+
+from fabricks.context import SPARK
+from fabricks.context.log import DEFAULT_LOGGER
+from fabricks.metastore.dbobject import DbObject
+from fabricks.utils._types import DataFrameLike
+
+
+class View(DbObject):
+    @staticmethod
+    def create_or_replace(
+        df: Union[DataFrame, pd.DataFrame],
+        *dependencies,
+        spark: Optional[SparkSession] = None,
+    ) -> str:
+        if spark is None:
+            if isinstance(df, DataFrameLike):
+                spark = df.sparkSession
+            else:
+                spark = SPARK
+
+        assert spark is not None
+
+        uuid = str(uuid4().hex)
+        df = spark.createDataFrame(df) if isinstance(df, pd.DataFrame) else df
+        if dependencies:
+            for d in dependencies:
+                df = df.join(d.where("1 == 2"), how="leftanti")
+
+        df.createOrReplaceGlobalTempView(uuid)
+        return uuid
+
+
+def create_or_replace_global_temp_view(
+    name: str,
+    df: DataFrame,
+    uuid: Optional[bool] = False,
+    job: Optional[Any] = None,
+) -> str:
+    if uuid:
+        name = f"{name}__{str(uuid4().hex)}"
+
+    if job is None:
+        job = name.split("__")[0]
+
+    DEFAULT_LOGGER.debug(f"create global temp view {name}", extra={"label": job})
+    df.createOrReplaceGlobalTempView(name)
+
+    return f"global_temp.{name}"

fabricks/utils/_types.py

@@ -0,0 +1,6 @@
+from typing import Union
+
+from pyspark.sql import DataFrame
+from pyspark.sql.connect.dataframe import DataFrame as ConnectDataFrame
+
+DataFrameLike = Union[DataFrame, ConnectDataFrame]

fabricks/utils/azure_queue.py

@@ -0,0 +1,93 @@
+import json
+from typing import TYPE_CHECKING, Optional, Union
+
+from azure.core.exceptions import ResourceExistsError
+from azure.storage.queue import QueueClient
+
+if TYPE_CHECKING:
+    from azure.core.credentials import TokenCredential
+
+
+class AzureQueue:
+    def __init__(
+        self,
+        name: str,
+        storage_account: Optional[str] = None,
+        access_key: Optional[str] = None,
+        connection_string: Optional[str] = None,
+        credential: "Optional[TokenCredential]" = None,
+    ):
+        self.name = name
+        self.storage_account = storage_account
+        if connection_string is None:
+            assert storage_account
+            assert access_key or credential, "Either access_key or credential must be provided"
+            self.storage_account = storage_account
+            self.access_key = access_key
+            self.credential = credential
+            connection_string = (
+                f"DefaultEndpointsProtocol=https;AccountName={self.storage_account};AccountKey={self.access_key};EndpointSuffix=core.windows.net"
+                if access_key
+                else None
+            )
+
+        assert connection_string
+        self.connection_string = connection_string
+        self._queue_client = None
+
+    @property
+    def queue_client(self) -> QueueClient:
+        if not self._queue_client:
+            if self.connection_string is not None:
+                self._queue_client = QueueClient.from_connection_string(self.connection_string, queue_name=self.name)
+            else:
+                assert self.storage_account and (self.access_key or self.credential), (
+                    "Either access_key or credential must be provided"
+                )
+                self._queue_client = QueueClient(
+                    account_url=f"https://{self.storage_account}.queue.core.windows.net",
+                    queue_name=self.name,
+                    credential=self.access_key if self.access_key else self.credential,
+                )
+        return self._queue_client
+
+    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()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args, **kwargs):
+        if self._queue_client is not None:
+            self._queue_client.close()

fabricks/utils/azure_table.py

@@ -0,0 +1,154 @@
+import time
+from typing import TYPE_CHECKING, List, Optional, Union
+
+from azure.data.tables import TableClient, TableServiceClient
+from pyspark.sql import DataFrame
+from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential
+
+from fabricks.utils._types import DataFrameLike
+
+if TYPE_CHECKING:
+    from azure.core.credentials import TokenCredential
+
+
+class AzureTable:
+    def __init__(
+        self,
+        name: str,
+        storage_account: Optional[str] = None,
+        access_key: Optional[str] = None,
+        connection_string: Optional[str] = None,
+        credential: "Optional[TokenCredential]" = None,
+    ):
+        self.name = name
+
+        if connection_string is None:
+            assert storage_account, "storage_account must be provided if connection_string is not set"
+            assert access_key or credential, "Either access_key or credential must be provided"
+            self.storage_account = storage_account
+            self.access_key = access_key
+            self.credential = credential
+            self.storage_account = storage_account
+
+            connection_string = (
+                f"DefaultEndpointsProtocol=https;AccountName={self.storage_account};AccountKey={self.access_key};EndpointSuffix=core.windows.net"
+                if access_key
+                else None
+            )
+
+        assert connection_string
+        self.connection_string = connection_string
+
+        self._table_client = None
+
+    @property
+    def table_service_client(self) -> TableServiceClient:
+        if not self._table_client:
+            if self.connection_string is None:
+                return TableServiceClient(
+                    endpoint=f"https://{self.storage_account}.table.core.windows.net",
+                    credential=self.credential,
+                )
+            self._table_client = TableServiceClient.from_connection_string(self.connection_string)
+        return self._table_client
+
+    @property
+    def table(self) -> TableClient:
+        return self.create_if_not_exists()
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=1, max=10),
+        retry=retry_if_exception_type((Exception)),
+        reraise=True,
+    )
+    def create_if_not_exists(self) -> TableClient:
+        return self.table_service_client.create_table_if_not_exists(table_name=self.name)
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=1, max=10),
+        retry=retry_if_exception_type((Exception)),
+        reraise=True,
+    )
+    def drop(self):
+        self.table_service_client.delete_table(self.name)
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=1, max=10),
+        retry=retry_if_exception_type((Exception)),
+        reraise=True,
+    )
+    def query(self, query: str) -> List:
+        return list(self.table.query_entities(query))
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=1, max=10),
+        retry=retry_if_exception_type((Exception)),
+        reraise=True,
+    )
+    def list_all(self) -> List:
+        return self.query("")
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args, **kwargs) -> None:
+        if self._table_client is not None:
+            self._table_client.close()
+
+    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, DataFrameLike):
+            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, DataFrameLike):
+            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/fdict.py

@@ -0,0 +1,240 @@
+from typing import Any, Callable, Dict, List, Optional, TypeVar, Union, overload
+
+T = TypeVar("T")
+
+
+class FDict:
+    """
+    A flexible dictionary wrapper that provides type-safe access to nested data structures
+    with convenient conversion methods.
+    """
+
+    def __init__(self, options: Union[Dict[str, Any], Any, None] = None):
+        """
+        Initialize FDict with a dictionary of options.
+
+        Args:
+            options: Input dictionary. If None, creates an empty dictionary.
+        """
+        self.options = options if options is not None else {}
+
+    def __getitem__(self, key: str) -> Any:
+        """Enable dictionary-like access with [] operator."""
+        return self.options[key]
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        """Enable dictionary-like value setting with [] operator."""
+        self.options[key] = value
+
+    def __contains__(self, key: str) -> bool:
+        """Enable 'in' operator for membership testing."""
+        return key in self.options
+
+    def __repr__(self) -> str:
+        """Return string representation of the FDict."""
+        return f"FDict({self.options})"
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert FDict to a regular dictionary."""
+        return self.options
+
+    @overload
+    def get(self, key: str) -> Optional[Any]: ...
+
+    @overload
+    def get(self, key: str, default: T) -> Union[Any, T]: ...
+
+    def get(self, key: str, default: Any = None) -> Any:
+        """
+        Get a value from the dictionary with an optional default.
+
+        Args:
+            key: The key to look up
+            default: Value to return if key is not found
+
+        Returns:
+            The value associated with the key or the default value
+        """
+        return self.options.get(key, default)
+
+    def get_list(self, key: str, default: Optional[List[Any]] = None) -> List[Any]:
+        """
+        Get a value as a list, converting single items to a single-item list.
+
+        Args:
+            key: The key to look up
+            default: Default value if key is not found
+
+        Returns:
+            A list containing the value(s)
+        """
+        values = self.options.get(key, default if default is not None else [])
+        if values is None:
+            return []
+
+        return [values] if not isinstance(values, list) else values
+
+    def get_boolean(self, key: str, default: Optional[bool] = None) -> Optional[bool]:
+        """
+        Get a value as a boolean, with string conversion support.
+
+        Args:
+            key: The key to look up
+            default: Default value if key is not found
+
+        Returns:
+            Boolean value of the key, or default if key not found
+        """
+        value = self.options.get(key)
+
+        if value is None:
+            return default
+        if isinstance(value, bool):
+            return value
+        if isinstance(value, str):
+            return value.lower() in ("true", "1", "yes", "on")
+
+        return bool(value)
+
+    def get_dict(self, key: str, default: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
+        """
+        Get a nested dictionary, with a default empty dict if not found.
+
+        Args:
+            key: The key to look up
+            default: Default value if key is not found
+
+        Returns:
+            Dictionary value of the key, or default if key not found
+        """
+        return self.options.get(key, default if default is not None else {})
+
+    def get_nested(self, *keys: str, default: Any = None) -> Any:
+        """
+        Access nested dictionary values using a sequence of keys.
+
+        Args:
+            *keys: Sequence of keys to traverse
+            default: Default value if path not found
+
+        Returns:
+            Value at the nested path, or default if path not found
+        """
+        current = self.options
+        for key in keys:
+            if not isinstance(current, dict):
+                return default
+            if key not in current:
+                return default
+            current = current[key]
+
+        return current
+
+    def set_nested(self, *keys: str, value: Any) -> None:
+        """
+        Set a value in a nested dictionary path, creating intermediate dictionaries as needed.
+
+        Args:
+            *keys: Sequence of keys defining the path
+            value: Value to set at the path
+        """
+        current = self.options
+        for key in keys[:-1]:
+            current = current.setdefault(key, {})
+
+        current[keys[-1]] = value
+
+    def filter(self, predicate: Callable[[str, Any], bool]) -> "FDict":
+        """
+        Create a new FDict with key-value pairs that satisfy the predicate function.
+
+        Args:
+            predicate: Lambda function that takes key and value as arguments and returns bool
+
+        Returns:
+            New FDict containing only the filtered key-value pairs
+
+        Example:
+            # Get all items with numeric values greater than 10
+            filtered = fdict.filter(lambda k, v: isinstance(v, (int, float)) and v > 10)
+        """
+        filtered_dict = {k: v for k, v in self.options.items() if predicate(k, v)}
+        return FDict(filtered_dict)
+
+    def filter_keys(self, predicate: Callable[[str], bool]) -> "FDict":
+        """
+        Create a new FDict with keys that satisfy the predicate function.
+
+        Args:
+            predicate: Lambda function that takes key as argument and returns bool
+
+        Returns:
+            New FDict containing only the filtered keys
+
+        Example:
+            # Get all items with keys starting with 'user_'
+            filtered = fdict.filter_keys(lambda k: k.startswith('user_'))
+        """
+        return self.filter(lambda k, _: predicate(k))
+
+    def filter_values(self, predicate: Callable[[Any], bool]) -> "FDict":
+        """
+        Create a new FDict with values that satisfy the predicate function.
+
+        Args:
+            predicate: Lambda function that takes value as argument and returns bool
+
+        Returns:
+            New FDict containing only the filtered values
+
+        Example:
+            # Get all items with string values
+            filtered = fdict.filter_values(lambda v: isinstance(v, str))
+        """
+        return self.filter(lambda _, v: predicate(v))
+
+    def map_values(self, transform: Callable[[Any], Any]) -> "FDict":
+        """
+        Create a new FDict with transformed values using the provided function.
+
+        Args:
+            transform: Lambda function that takes a value and returns transformed value
+
+        Returns:
+            New FDict containing transformed values
+
+        Example:
+            # Convert all string values to uppercase
+            transformed = fdict.map_values(lambda v: v.upper() if isinstance(v, str) else v)
+        """
+        transformed_dict = {k: transform(v) for k, v in self.options.items()}
+        return FDict(transformed_dict)
+
+    def deep_filter(self, predicate: Callable[[str, Any], bool]) -> "FDict":
+        """
+        Recursively filter nested dictionaries using the predicate function.
+
+        Args:
+            predicate: Lambda function that takes key and value as arguments and returns bool
+
+        Returns:
+            New FDict with filtered nested structure
+
+        Example:
+            # Filter all nested numeric values greater than 10
+            filtered = fdict.deep_filter(lambda k, v:
+                not isinstance(v, dict) and isinstance(v, (int, float)) and v > 10)
+        """
+
+        def filter_recursive(d: Dict[str, Any]) -> Dict[str, Any]:
+            result = {}
+            for k, v in d.items():
+                if isinstance(v, dict):
+                    filtered_nested = filter_recursive(v)
+                    if filtered_nested:  # Only include non-empty nested dicts
+                        result[k] = filtered_nested
+                elif predicate(k, v):
+                    result[k] = v
+            return result
+
+        return FDict(filter_recursive(self.options))