squirrels 0.1.0__py3-none-any.whl → 0.6.0.post0__py3-none-any.whl

squirrels/_utils.py

@@ -1,149 +1,540 @@
-from typing import List, Dict, Optional, Union, Any
-from types import ModuleType
-from pathlib import Path
-from importlib.machinery import SourceFileLoader
-import json
-
-from squirrels._timed_imports import jinja2 as j2, pandas as pd, pd_types
-
-FilePath = Union[str, Path]
-
-
-# Custom Exceptions
-class InvalidInputError(Exception):
-    pass
-
-class ConfigurationError(Exception):
-    pass
-
-class AbstractMethodCallError(NotImplementedError):
-    def __init__(self, cls, method, more_message = ""):
-        message = f"Abstract method {method}() not implemented in {cls.__name__}."
-        super().__init__(message + more_message)
-
-
-# Utility functions/variables
-j2_env = j2.Environment(loader=j2.FileSystemLoader('.'))
-
-
-def import_file_as_module(filepath: Optional[FilePath]) -> ModuleType:
-    """
-    Imports a python file as a module.
-
-    Parameters:
-        filepath: The path to the file to import.
-
-    Returns:
-        The imported module.
-    """
-    filepath = str(filepath) if filepath is not None else None
-    return SourceFileLoader(filepath, filepath).load_module() if filepath is not None else None
-
-
-def join_paths(*paths: FilePath) -> Path:
-    """
-    Joins paths together.
-
-    Parameters:
-        paths: The paths to join.
-
-    Returns:
-        The joined path.
-    """
-    return Path(*paths)
-
-
-def normalize_name(name: str) -> str:
-    """
-    Normalizes names to the convention of the squirrels manifest file.
-
-    Parameters:
-        name: The name to normalize.
-
-    Returns:
-        The normalized name.
-    """
-    return name.replace('-', '_')
-
-
-def normalize_name_for_api(name: str) -> str:
-    """
-    Normalizes names to the REST API convention.
-
-    Parameters:
-        name: The name to normalize.
-
-    Returns:
-        The normalized name.
-    """
-    return name.replace('_', '-')
-
-
-def get_row_value(row: pd.Series, value: str) -> Any:
-    """
-    Gets the value of a row from a pandas Series.
-
-    Parameters:
-        row: The row to get the value from.
-        value: The name of the column to get the value from.
-
-    Returns:
-        The value of the column.
-
-    Raises:
-        ConfigurationError: If the column does not exist.
-    """
-    try:
-        result = row[value]
-    except KeyError as e:
-        raise ConfigurationError(f'Column name "{value}" does not exist') from e
-    return result
-
-
-def df_to_json(df: pd.DataFrame, dimensions: List[str] = None) -> Dict[str, Any]:
-    """
-    Convert a pandas DataFrame to the same JSON format that the dataset result API of Squirrels outputs.
-
-    Parameters:
-        df: The dataframe to convert into JSON
-        dimensions: The list of declared dimensions. If None, all non-numeric columns are assumed as dimensions
-
-    Returns:
-        The JSON response of a Squirrels dataset result API
-    """
-    in_df_json = json.loads(df.to_json(orient='table', index=False))
-    out_fields = []
-    non_numeric_fields = []
-    for in_column in in_df_json["schema"]["fields"]:
-        col_name: str = in_column["name"]
-        out_column = {"name": col_name, "type": in_column["type"]}
-        out_fields.append(out_column)
-        
-        if not pd_types.is_numeric_dtype(df[col_name].dtype):
-            non_numeric_fields.append(col_name)
-    
-    out_dimensions = non_numeric_fields if dimensions is None else dimensions
-    out_schema = {"fields": out_fields, "dimensions": out_dimensions}
-    return {"response_version": 0, "schema": out_schema, "data": in_df_json["data"]}
-
-
-def load_json_or_comma_delimited_str_as_list(input_str: str) -> List[str]:
-    """
-    Given a string, load it as a list either by json string or comma delimited value
-
-    Parameters:
-        input_str: The input string
-    
-    Returns:
-        The list representation of the input string
-    """
-    output = None
-    try:
-        output = json.loads(input_str)
-    except json.decoder.JSONDecodeError:
-        pass
-    
-    if isinstance(output, list):
-        return output
-    else:
-        return [] if input_str == "" else input_str.split(",")
+from typing import Sequence, Optional, Union, TypeVar, Callable, Iterable, Literal, Any
+from datetime import datetime
+from pathlib import Path
+import os, time, logging, json, duckdb, polars as pl, yaml
+import jinja2 as j2, jinja2.nodes as j2_nodes
+import sqlglot, sqlglot.expressions, asyncio, hashlib, inspect, base64
+
+from . import _constants as c
+from ._exceptions import ConfigurationError
+
+FilePath = Union[str, Path]
+
+# Polars <-> Squirrels dtypes mappings (except Decimal)
+polars_dtypes_to_sqrl_dtypes: dict[type[pl.DataType], list[str]] = {
+    pl.String: ["string", "varchar", "char", "text"],
+    pl.Int8: ["tinyint", "int1"],
+    pl.Int16: ["smallint", "short", "int2"],
+    pl.Int32: ["integer", "int", "int4"],
+    pl.Int64: ["bigint", "long", "int8"],
+    pl.Float32: ["float", "float4", "real"],
+    pl.Float64: ["double", "float8"],
+    pl.Boolean: ["boolean", "bool", "logical"],
+    pl.Date: ["date"],
+    pl.Time: ["time"],
+    pl.Datetime: ["timestamp", "datetime"],
+    pl.Duration: ["interval"],
+    pl.Binary: ["blob", "binary", "varbinary"]
+}
+
+sqrl_dtypes_to_polars_dtypes: dict[str, type[pl.DataType]] = {
+    sqrl_type: k for k, v in polars_dtypes_to_sqrl_dtypes.items() for sqrl_type in v
+}
+
+
+## Other utility classes
+
+class Logger(logging.Logger):
+    def info(self, msg: str, *, data: dict[str, Any] = {}, **kwargs) -> None:
+        super().info(msg, extra={"data": data}, **kwargs)
+
+    def log_activity_time(self, activity: str, start_timestamp: float, *, additional_data: dict[str, Any] = {}) -> None:
+        end_timestamp = time.time()
+        time_taken = round((end_timestamp-start_timestamp) * 10**3, 3)
+        data = {
+            "activity": activity,
+            "start_timestamp": start_timestamp, 
+            "end_timestamp": end_timestamp, 
+            "time_taken_ms": time_taken,
+            **additional_data
+        }
+        self.info(f'Time taken for "{activity}": {time_taken}ms', data=data)
+
+
+class EnvironmentWithMacros(j2.Environment):
+    def __init__(self, logger: logging.Logger, loader: j2.FileSystemLoader, *args, **kwargs):
+        super().__init__(*args, loader=loader, **kwargs)
+        self._logger = logger
+        self._macros = self._load_macro_templates(logger)
+
+    def _load_macro_templates(self, logger: logging.Logger) -> str:
+        macros_dirs = self._get_macro_folders_from_packages()
+        macro_templates = []
+        for macros_dir in macros_dirs:
+            for root, _, files in os.walk(macros_dir):
+                files: list[str]
+                for filename in files:
+                    if any(filename.endswith(x) for x in [".sql", ".j2", ".jinja", ".jinja2"]):
+                        filepath = Path(root, filename)
+                        logger.info(f"Loaded macros from: {filepath}")
+                        with open(filepath, 'r') as f:
+                            content = f.read()
+                        macro_templates.append(content)
+        return '\n'.join(macro_templates)
+    
+    def _get_macro_folders_from_packages(self) -> list[Path]:
+        assert isinstance(self.loader, j2.FileSystemLoader)
+        packages_folder = Path(self.loader.searchpath[0], c.PACKAGES_FOLDER)
+        
+        subdirectories = []
+        if os.path.exists(packages_folder):
+            for item in os.listdir(packages_folder):
+                item_path = Path(packages_folder, item)
+                if os.path.isdir(item_path):
+                    subdirectories.append(Path(item_path, c.MACROS_FOLDER))
+        
+        subdirectories.append(Path(self.loader.searchpath[0], c.MACROS_FOLDER))
+        return subdirectories
+
+    def _parse(self, source: str, name: str | None, filename: str | None) -> j2_nodes.Template:
+        source = self._macros + source
+        return super()._parse(source, name, filename)
+
+
+## Utility functions/variables
+
+def render_string(raw_str: str, *, project_path: str = ".", **kwargs) -> str:
+    """
+    Given a template string, render it with the given keyword arguments
+
+    Arguments:
+        raw_str: The template string
+        kwargs: The keyword arguments
+
+    Returns:
+        The rendered string
+    """
+    j2_env = j2.Environment(loader=j2.FileSystemLoader(project_path))
+    template = j2_env.from_string(raw_str)
+    return template.render(kwargs)
+
+
+def read_file(filepath: FilePath) -> str:
+    """
+    Reads a file and return its content if required
+
+    Arguments:
+        filepath (str | pathlib.Path): The path to the file to read
+
+    Returns:
+        Content of the file, or None if doesn't exist and not required
+    """
+    try:
+        with open(filepath, 'r') as f:
+            return f.read()
+    except FileNotFoundError as e:
+        raise ConfigurationError(f"Required file not found: '{str(filepath)}'") from e
+
+
+def normalize_name(name: str) -> str:
+    """
+    Normalizes names to the convention of the squirrels manifest file (with underscores instead of dashes).
+
+    Arguments:
+        name: The name to normalize.
+
+    Returns:
+        The normalized name.
+    """
+    return name.replace('-', '_')
+
+
+def normalize_name_for_api(name: str) -> str:
+    """
+    Normalizes names to the REST API convention (with dashes instead of underscores).
+
+    Arguments:
+        name: The name to normalize.
+
+    Returns:
+        The normalized name.
+    """
+    return name.replace('_', '-')
+
+
+def load_json_or_comma_delimited_str_as_list(input_str: Union[str, Sequence]) -> Sequence[str]:
+    """
+    Given a string, load it as a list either by json string or comma delimited value
+
+    Arguments:
+        input_str: The input string
+    
+    Returns:
+        The list representation of the input string
+    """
+    if not isinstance(input_str, str):
+        return (input_str)
+    
+    output = None
+    try:
+        output = json.loads(input_str)
+    except json.decoder.JSONDecodeError:
+        pass
+    
+    if isinstance(output, list):
+        return output
+    elif input_str == "":
+        return []
+    else:
+        return [x.strip() for x in input_str.split(",")]
+
+
+X = TypeVar('X'); Y = TypeVar('Y')
+def process_if_not_none(input_val: Optional[X], processor: Callable[[X], Y]) -> Optional[Y]:
+    """
+    Given a input value and a function that processes the value, return the output of the function unless input is None
+
+    Arguments:
+        input_val: The input value
+        processor: The function that processes the input value
+    
+    Returns:
+        The output type of "processor" or None if input value if None
+    """
+    if input_val is None:
+        return None
+    return processor(input_val)
+
+
+def _read_duckdb_init_sql(
+    *,
+    datalake_db_path: str | None = None,
+) -> str:
+    """
+    Reads and caches the duckdb init file content.
+    Returns None if file doesn't exist or is empty.
+    """
+    try:
+        init_contents = []
+        global_init_path = Path(os.path.expanduser('~'), c.GLOBAL_ENV_FOLDER, c.DUCKDB_INIT_FILE)
+        if global_init_path.exists():
+            with open(global_init_path, 'r') as f:
+                init_contents.append(f.read())
+        
+        if Path(c.DUCKDB_INIT_FILE).exists():
+            with open(c.DUCKDB_INIT_FILE, 'r') as f:
+                init_contents.append(f.read())
+
+        if datalake_db_path:
+            attach_stmt = f"ATTACH '{datalake_db_path}' AS vdl (READ_ONLY);"
+            init_contents.append(attach_stmt)
+            use_stmt = f"USE vdl;"
+            init_contents.append(use_stmt)
+                
+        init_sql = "\n\n".join(init_contents).strip()
+        return init_sql
+    except Exception as e:
+        raise ConfigurationError(f"Failed to read {c.DUCKDB_INIT_FILE}: {str(e)}") from e
+
+def create_duckdb_connection(
+    db_path: str | Path = ":memory:", 
+    *, 
+    datalake_db_path: str | None = None
+) -> duckdb.DuckDBPyConnection:
+    """
+    Creates a DuckDB connection and initializes it with statements from duckdb init file
+
+    Arguments:
+        filepath: Path to the DuckDB database file. Defaults to in-memory database.
+        datalake_db_path: The path to the VDL catalog database if applicable. If exists, this is attached as 'vdl' (READ_ONLY). Default is None.
+    
+    Returns:
+        A DuckDB connection (which must be closed after use)
+    """
+    conn = duckdb.connect(db_path)
+    
+    try:
+        init_sql = _read_duckdb_init_sql(datalake_db_path=datalake_db_path)
+        conn.execute(init_sql)
+    except Exception as e:
+        conn.close()
+        raise ConfigurationError(f"Failed to execute {c.DUCKDB_INIT_FILE}: {str(e)}") from e
+    
+    return conn
+
+
+def run_sql_on_dataframes(sql_query: str, dataframes: dict[str, pl.LazyFrame]) -> pl.DataFrame:
+    """
+    Runs a SQL query against a collection of dataframes
+
+    Arguments:
+        sql_query: The SQL query to run
+        dataframes: A dictionary of table names to their polars LazyFrame
+    
+    Returns:
+        The result as a polars Dataframe from running the query
+    """
+    duckdb_conn = create_duckdb_connection()
+    
+    try:
+        for name, df in dataframes.items():
+            duckdb_conn.register(name, df)
+        
+        result_df = duckdb_conn.sql(sql_query).pl()
+    finally:
+        duckdb_conn.close()
+    
+    return result_df
+
+
+async def run_polars_sql_on_dataframes(
+    sql_query: str, dataframes: dict[str, pl.LazyFrame], *, timeout_seconds: float = 2.0, max_rows: int | None = None
+) -> pl.DataFrame:
+    """
+    Runs a SQL query against a collection of dataframes using Polars SQL (more secure than DuckDB for user input).
+    
+    Arguments:
+        sql_query: The SQL query to run (Polars SQL dialect)
+        dataframes: A dictionary of table names to their polars LazyFrame
+        timeout_seconds: Maximum execution time in seconds (default 2.0)
+        max_rows: Maximum number of rows to collect. Collects at most max_rows + 1 rows
+                  to allow overflow detection without loading unbounded results into memory.
+    
+    Returns:
+        The result as a polars DataFrame from running the query (limited to max_rows + 1)
+    
+    Raises:
+        ConfigurationError: If the query is invalid or insecure
+    """
+    # Validate the SQL query
+    _validate_sql_query_security(sql_query, dataframes)
+    
+    # Execute with timeout
+    try:
+        loop = asyncio.get_event_loop()
+        result = await asyncio.wait_for(
+            loop.run_in_executor(None, _run_polars_sql_sync, sql_query, dataframes, max_rows),
+            timeout=timeout_seconds
+        )
+        return result
+    except asyncio.TimeoutError as e:
+        raise ConfigurationError(f"SQL query execution exceeded timeout of {timeout_seconds} seconds") from e
+
+
+def _run_polars_sql_sync(sql_query: str, dataframes: dict[str, pl.LazyFrame], max_rows: int | None) -> pl.DataFrame:
+    """
+    Synchronous execution of Polars SQL.
+    
+    Arguments:
+        sql_query: The SQL query to run
+        dataframes: A dictionary of table names to their polars LazyFrame
+        max_rows: Maximum number of rows to collect.
+    """
+    ctx = pl.SQLContext(**dataframes)
+    result = ctx.execute(sql_query, eager=False)
+    if max_rows is not None:
+        result = result.limit(max_rows)
+    return result.collect()
+
+
+def _validate_sql_query_security(sql_query: str, dataframes: dict[str, pl.LazyFrame]) -> None:
+    """
+    Validates that a SQL query is safe to execute.
+    
+    Enforces:
+    - Single statement only
+    - Read-only operations (SELECT/WITH/UNION)
+    - Table references limited to registered frames (excluding CTE names)
+    
+    Arguments:
+        sql_query: The SQL query to validate
+        dataframes: Dictionary of allowed table names
+        
+    Raises:
+        ConfigurationError: If validation fails
+    """
+    try:
+        parsed = sqlglot.parse(sql_query)
+    except Exception as e:
+        raise ConfigurationError(f"Failed to parse SQL query: {str(e)}") from e
+    
+    # Enforce single statement
+    if len(parsed) != 1:
+        raise ConfigurationError(f"Only single SQL statements are allowed. Found {len(parsed)} statements.")
+    
+    statement = parsed[0]
+    
+    # Enforce read-only: allow SELECT, WITH (CTE), UNION, INTERSECT, EXCEPT
+    allowed_types = (
+        sqlglot.expressions.Select,
+        sqlglot.expressions.Union,
+        sqlglot.expressions.Intersect,
+        sqlglot.expressions.Except,
+    )
+    
+    if not isinstance(statement, allowed_types):
+        raise ConfigurationError(
+            f"Only read-only SQL statements (SELECT, WITH, UNION, INTERSECT, EXCEPT) are allowed. "
+            f"Found: {type(statement).__name__}"
+        )
+    
+    # Collect CTE names (these are temporary tables created by WITH clauses)
+    cte_names: set[str] = set()
+    for cte in statement.find_all(sqlglot.expressions.CTE):
+        if cte.alias:
+            cte_names.add(cte.alias)
+    
+    # Validate table references (excluding CTE names)
+    allowed_tables = set(dataframes.keys()) | cte_names
+    for table in statement.find_all(sqlglot.expressions.Table):
+        table_name = table.name
+        if table_name not in allowed_tables:
+            raise ConfigurationError(
+                f"Table reference '{table_name}' is not allowed. "
+                f"Only the following tables are available: {sorted(dataframes.keys())}"
+            )
+
+
+def load_yaml_config(filepath: FilePath) -> dict:
+    """
+    Loads a YAML config file
+
+    Arguments:
+        filepath: The path to the YAML file
+    
+    Returns:
+        A dictionary representation of the YAML file
+    """
+    try:
+        with open(filepath, 'r') as f:
+            content = yaml.safe_load(f)
+            content = content if content else {}
+        
+        if not isinstance(content, dict):
+            raise yaml.YAMLError(f"Parsed content from YAML file must be a dictionary. Got: {content}")
+        
+        return content
+    except yaml.YAMLError as e:
+        raise ConfigurationError(f"Failed to parse yaml file: {filepath}") from e
+
+
+def run_duckdb_stmt(
+    logger: Logger, duckdb_conn: duckdb.DuckDBPyConnection, stmt: str, *, params: dict[str, Any] | None = None, 
+    model_name: str | None = None, redacted_values: list[str] = []
+) -> duckdb.DuckDBPyConnection:
+    """
+    Runs a statement on a DuckDB connection
+
+    Arguments:
+        logger: The logger to use
+        duckdb_conn: The DuckDB connection
+        stmt: The statement to run
+        params: The parameters to use
+        redacted_values: The values to redact
+    """
+    redacted_stmt = stmt
+    for value in redacted_values:
+        redacted_stmt = redacted_stmt.replace(value, "[REDACTED]")
+    
+    for_model_name = f" for model '{model_name}'" if model_name is not None else ""
+    logger.debug(f"Running SQL statement{for_model_name}:\n{redacted_stmt}")
+    try:
+        return duckdb_conn.execute(stmt, params)
+    except duckdb.ParserException as e:
+        logger.error(f"Failed to run statement: {redacted_stmt}", exc_info=e)
+        raise e
+
+
+def get_current_time() -> str:
+    """
+    Returns the current time in the format HH:MM:SS.ms
+    """
+    return datetime.now().strftime('%H:%M:%S.%f')[:-3]
+
+
+def parse_dependent_tables(sql_query: str, all_table_names: Iterable[str]) -> tuple[set[str], sqlglot.Expression]:
+    """
+    Parses the dependent tables from a SQL query
+
+    Arguments:
+        sql_query: The SQL query to parse
+        all_table_names: The list of all table names
+    
+    Returns:
+        The set of dependent tables
+    """
+    # Parse the SQL query and extract all table references
+    parsed = sqlglot.parse_one(sql_query)
+    dependencies = set()
+    
+    # Collect all table references from the parsed SQL
+    for table in parsed.find_all(sqlglot.expressions.Table):
+        if table.name in set(all_table_names):
+            dependencies.add(table.name)
+    
+    return dependencies, parsed
+
+
+async def asyncio_gather(coroutines: list):
+    tasks = [asyncio.create_task(coro) for coro in coroutines]
+    
+    try:
+        return await asyncio.gather(*tasks)
+    except BaseException:
+        # Cancel all tasks
+        for task in tasks:
+            if not task.done():
+                task.cancel()
+        # Wait for tasks to be cancelled
+        await asyncio.gather(*tasks, return_exceptions=True)
+        raise
+
+
+def hash_string(input_str: str, salt: str) -> str:
+    """
+    Hashes a string using SHA-256
+    """
+    return hashlib.sha256((input_str + salt).encode()).hexdigest()
+
+
+T = TypeVar('T')
+def call_func(func: Callable[..., T], **kwargs) -> T:
+    """
+    Calls a function with the given arguments if func expects arguments, otherwise calls func without arguments
+    """
+    sig = inspect.signature(func)
+    # Filter kwargs to only include parameters that the function accepts
+    filtered_kwargs = {k: v for k, v in kwargs.items() if k in sig.parameters}
+    return func(**filtered_kwargs)
+
+
+def generate_pkce_challenge(code_verifier: str) -> str:
+    """Generate PKCE code challenge from code verifier"""
+    # Generate SHA256 hash of code_verifier
+    verifier_hash = hashlib.sha256(code_verifier.encode('utf-8')).digest()
+    # Base64 URL encode (without padding)
+    expected_challenge = base64.urlsafe_b64encode(verifier_hash).decode('utf-8').rstrip('=')
+    return expected_challenge
+
+
+def to_title_case(input_str: str) -> str:
+    """Convert a string to title case"""
+    spaced_str = input_str.replace('_', ' ').replace('-', ' ')
+    return spaced_str.title()
+
+
+def to_bool(val: object) -> bool:
+    """Convert common truthy/falsey representations to a boolean.
+
+    Accepted truthy values (case-insensitive): "1", "true", "t", "yes", "y", "on".
+    All other values are considered falsey. None is falsey.
+    """
+    if isinstance(val, bool):
+        return val
+    if val is None:
+        return False
+    s = str(val).strip().lower()
+    return s in ("1", "true", "t", "yes", "y", "on")
+
+
+ACCESS_LEVEL = Literal["admin", "member", "guest"]
+
+def get_access_level_rank(access_level: ACCESS_LEVEL) -> int:
+    """Get the rank of an access level. Lower ranks have more privileges."""
+    return { "admin": 1, "member": 2, "guest": 3 }.get(access_level.lower(), 1)
+
+def user_has_elevated_privileges(user_access_level: ACCESS_LEVEL, required_access_level: ACCESS_LEVEL) -> bool:
+    """Check if a user has privilege to access a resource"""
+    user_access_level_rank = get_access_level_rank(user_access_level)
+    required_access_level_rank = get_access_level_rank(required_access_level)
+    return user_access_level_rank <= required_access_level_rank