squirrels 0.4.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

squirrels/_sources.py

@@ -0,0 +1,110 @@
+from typing import Any
+from pydantic import BaseModel, Field, model_validator
+import time, sqlglot, yaml
+
+from . import _utils as u, _constants as c, _model_configs as mc
+
+
+class UpdateHints(BaseModel):
+    increasing_column: str | None = Field(default=None)
+    strictly_increasing: bool = Field(default=True, description="Delete the max value of the increasing column, ignored if selective_overwrite_value is set")
+    selective_overwrite_value: Any = Field(default=None, description="Delete all values of the increasing column greater than or equal to this value")
+
+
+class Source(mc.ConnectionInterface, mc.ModelConfig):
+    table: str | None = Field(default=None)
+    load_to_vdl: bool = Field(default=False, description="Whether to load the data to the 'virtual data lake' (VDL)")
+    primary_key: list[str] = Field(default_factory=list)
+    update_hints: UpdateHints = Field(default_factory=UpdateHints)
+
+    def finalize_table(self, source_name: str):
+        if self.table is None:
+            self.table = source_name
+        return self
+    
+    def get_table(self) -> str:
+        assert self.table is not None, "Table must be set"
+        return self.table
+    
+    def get_cols_for_create_table_stmt(self) -> str:
+        cols_clause = ", ".join([f"{col.name} {col.type}" for col in self.columns])
+        return cols_clause
+    
+    def get_max_incr_col_query(self, source_name: str) -> str:
+        return f"SELECT max({self.update_hints.increasing_column}) FROM {source_name}"
+    
+    def get_query_for_upsert(self, dialect: str, conn_name: str, table_name: str, max_value_of_increasing_col: Any | None, *, full_refresh: bool = True) -> str:
+        select_cols = ", ".join([col.name for col in self.columns])
+        if full_refresh or max_value_of_increasing_col is None:
+            return f"SELECT {select_cols} FROM db_{conn_name}.{table_name}"
+        
+        increasing_col = self.update_hints.increasing_column
+        increasing_col_type = next(col.type for col in self.columns if col.name == increasing_col)
+        where_cond = f"{increasing_col}::{increasing_col_type} > '{max_value_of_increasing_col}'::{increasing_col_type}"
+        
+        # TODO: figure out if using pushdown query is worth it
+        # if dialect in ['postgres', 'mysql']:
+        #     pushdown_query = f"SELECT {select_cols} FROM {table_name} WHERE {where_cond}"
+        #     transpiled_query = sqlglot.transpile(pushdown_query, read='duckdb', write=dialect)[0].replace("'", "''")
+        #     return f"FROM {dialect}_query('db_{conn_name}', '{transpiled_query}')"
+        
+        return f"SELECT {select_cols} FROM db_{conn_name}.{table_name} WHERE {where_cond}"
+    
+
+class Sources(BaseModel):
+    sources: dict[str, Source] = Field(default_factory=dict)
+    
+    @model_validator(mode="before")
+    @classmethod
+    def convert_sources_list_to_dict(cls, data: dict[str, Any]) -> dict[str, Any]:
+        if "sources" in data and isinstance(data["sources"], list):
+            # Convert list of sources to dictionary
+            sources_dict = {}
+            for source in data["sources"]:
+                if isinstance(source, dict) and "name" in source:
+                    name = source.pop("name")  # Remove name from source config
+                    if name in sources_dict:
+                        raise u.ConfigurationError(f"Duplicate source name found: {name}")
+                    sources_dict[name] = source
+                else:
+                    raise u.ConfigurationError(f"All sources must have a name field in sources file")
+            data["sources"] = sources_dict
+        return data
+    
+    @model_validator(mode="after")
+    def validate_column_types(self):
+        for source_name, source in self.sources.items():
+            for col in source.columns:
+                if not col.type:
+                    raise u.ConfigurationError(f"Column '{col.name}' in source '{source_name}' must have a type specified")
+        return self
+    
+    def finalize_null_fields(self, env_vars: dict[str, str]):
+        for source_name, source in self.sources.items():
+            source.finalize_connection(env_vars)
+            source.finalize_table(source_name)
+        return self
+
+
+class SourcesIO:
+    @classmethod
+    def load_file(cls, logger: u.Logger, base_path: str, env_vars: dict[str, str]) -> Sources:
+        start = time.time()
+        
+        sources_path = u.Path(base_path, c.MODELS_FOLDER, c.SOURCES_FILE)
+        if sources_path.exists():
+            raw_content = u.read_file(sources_path)
+            rendered = u.render_string(raw_content, base_path=base_path, env_vars=env_vars)
+            sources_data = yaml.safe_load(rendered) or {}
+        else:
+            sources_data = {}
+        
+        if not isinstance(sources_data, dict):
+            raise u.ConfigurationError(
+                f"Parsed content from YAML file must be a dictionary. Got: {sources_data}"
+            )
+        
+        sources = Sources(**sources_data).finalize_null_fields(env_vars)
+        
+        logger.log_activity_time("loading sources", start)
+        return sources

squirrels/_utils.py

@@ -1,35 +1,33 @@
-from typing import Sequence, Optional, Union, TypeVar, Callable
-from pathlib import Path
-from pandas.api import types as pd_types
+from typing import Sequence, Optional, Union, TypeVar, Callable, Iterable, Literal, Any
 from datetime import datetime
-import os, time, logging, json, sqlite3, pandas as pd
+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]
 
-
-## Custom Exceptions
-
-class InvalidInputError(Exception):
-    """
-    Use this exception when the error is due to providing invalid inputs to the REST API
-    """
-    pass
-
-class ConfigurationError(Exception):
-    """
-    Use this exception when the server error is due to errors in the squirrels project instead of the squirrels framework/library
-    """
-    pass
-
-class FileExecutionError(Exception):
-    def __init__(self, message: str, error: Exception, *args) -> None:
-        t = "  "
-        new_message = f"\n" + message + f"\n{t}Produced error message:\n{t}{t}{error} (see above for more details on handled exception)"
-        super().__init__(new_message, *args)
-        self.error = error
+# Polars
+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", "decimal"], # Note: Polars Decimal type is considered unstable, so we use Float64 for "decimal"
+    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
@@ -40,7 +38,7 @@ class Logger(logging.Logger):
         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 }
         info = { "request_id": request_id } if request_id else {}
-        self.debug(f'Time taken for "{activity}": {time_taken}ms', extra={"data": data, "info": info})
+        self.info(f'Time taken for "{activity}": {time_taken}ms', extra={"data": data, "info": info})
 
 
 class EnvironmentWithMacros(j2.Environment):
@@ -85,14 +83,6 @@ class EnvironmentWithMacros(j2.Environment):
 
 ## Utility functions/variables
 
-def log_activity_time(logger: logging.Logger, activity: str, start_timestamp: float, *, request_id: str | None = None) -> 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 }
-    info = { "request_id": request_id } if request_id else {}
-    logger.debug(f'Time taken for "{activity}": {time_taken}ms', extra={"data": data, "info": info})
-
-
 def render_string(raw_str: str, *, base_path: str = ".", **kwargs) -> str:
     """
     Given a template string, render it with the given keyword arguments
@@ -115,7 +105,6 @@ def read_file(filepath: FilePath) -> str:
 
     Arguments:
         filepath (str | pathlib.Path): The path to the file to read
-        is_required: If true, throw error if file doesn't exist
 
     Returns:
         Content of the file, or None if doesn't exist and not required
@@ -129,7 +118,7 @@ def read_file(filepath: FilePath) -> str:
 
 def normalize_name(name: str) -> str:
     """
-    Normalizes names to the convention of the squirrels manifest file.
+    Normalizes names to the convention of the squirrels manifest file (with underscores instead of dashes).
 
     Arguments:
         name: The name to normalize.
@@ -142,7 +131,7 @@ def normalize_name(name: str) -> str:
 
 def normalize_name_for_api(name: str) -> str:
     """
-    Normalizes names to the REST API convention.
+    Normalizes names to the REST API convention (with dashes instead of underscores).
 
     Arguments:
         name: The name to normalize.
@@ -180,7 +169,7 @@ def load_json_or_comma_delimited_str_as_list(input_str: Union[str, Sequence]) ->
         return [x.strip() for x in input_str.split(",")]
 
 
-X, Y = TypeVar('X'), TypeVar('Y')
+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
@@ -197,60 +186,246 @@ def process_if_not_none(input_val: Optional[X], processor: Callable[[X], Y]) ->
     return processor(input_val)
 
 
-def run_sql_on_dataframes(sql_query: str, dataframes: dict[str, pd.DataFrame], do_use_duckdb: bool) -> pd.DataFrame:
+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)
+                
+        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 pandas Dataframe
+        dataframes: A dictionary of table names to their polars LazyFrame
     
     Returns:
-        The result as a pandas Dataframe from running the query
+        The result as a polars Dataframe from running the query
     """
-    if do_use_duckdb:
-        import duckdb
-        duckdb_conn = duckdb.connect()
-    else:
-        conn = sqlite3.connect(":memory:")
+    duckdb_conn = create_duckdb_connection()
     
     try:
         for name, df in dataframes.items():
-            if do_use_duckdb:
-                duckdb_conn.execute(f"CREATE TABLE {name} AS FROM df")
-            else:
-                df.to_sql(name, conn, index=False)
+            duckdb_conn.register(name, df)
         
-        return duckdb_conn.execute(sql_query).df() if do_use_duckdb else pd.read_sql(sql_query, conn)
+        result_df = duckdb_conn.sql(sql_query).pl()
     finally:
-        duckdb_conn.close() if do_use_duckdb else conn.close()
+        duckdb_conn.close()
+    
+    return result_df
 
 
-def df_to_json0(df: pd.DataFrame, dimensions: list[str] | None = None) -> dict:
+def load_yaml_config(filepath: FilePath) -> dict:
     """
-    Convert a pandas DataFrame to the response format that the dataset result API of Squirrels outputs.
+    Loads a YAML config file
 
     Arguments:
-        df: The dataframe to convert into an API response
-        dimensions: The list of declared dimensions. If None, all non-numeric columns are assumed as dimensions
-
+        filepath: The path to the YAML file
+    
     Returns:
-        The response of a Squirrels dataset result API
+        A dictionary representation of the YAML file
     """
-    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)
+    try:
+        with open(filepath, 'r') as f:
+            content = yaml.safe_load(f)
+            content = content if content else {}
         
-        if not pd_types.is_numeric_dtype(df[col_name].dtype):
-            non_numeric_fields.append(col_name)
+        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]")
     
-    out_dimensions = non_numeric_fields if dimensions is None else dimensions
-    dataset_json = {
-        "schema": { "fields": out_fields, "dimensions": out_dimensions },
-        "data": in_df_json["data"]
-    }
-    return dataset_json
+    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}", extra={"data": {"params": params}})
+    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 validate_pkce_challenge(code_verifier: str, code_challenge: str) -> bool:
+    """Validate PKCE code verifier against code challenge"""
+    # Generate expected challenge
+    expected_challenge = generate_pkce_challenge(code_verifier)
+    return expected_challenge == code_challenge
+
+
+def get_scheme(hostname: str | None) -> str:
+    """Get the scheme of the request"""
+    return "http" if hostname in ("localhost", "127.0.0.1") else "https"
+
+
+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

squirrels/_version.py

@@ -1,3 +1,3 @@
-__version__ = '0.4.0'
+__version__ = '0.5.0'
 
 sq_major_version, sq_minor_version, sq_patch_version = __version__.split('.')[:3]

squirrels/arguments.py

@@ -0,0 +1,7 @@
+from ._arguments.init_time_args import ConnectionsArgs, AuthProviderArgs, ParametersArgs, BuildModelArgs
+from ._arguments.run_time_args import ContextArgs, ModelArgs, DashboardArgs
+
+__all__ = [
+    "ConnectionsArgs", "AuthProviderArgs", "ParametersArgs", "BuildModelArgs", 
+    "ContextArgs", "ModelArgs", "DashboardArgs"
+]

squirrels/auth.py

@@ -0,0 +1,4 @@
+from ._schemas.auth_models import CustomUserFields, RegisteredUser
+from ._auth import ProviderConfigs, provider
+
+__all__ = ["CustomUserFields", "RegisteredUser", "ProviderConfigs", "provider"]

squirrels/connections.py

@@ -0,0 +1,3 @@
+from ._manifest import ConnectionProperties, ConnectionTypeEnum
+
+__all__ = ["ConnectionProperties", "ConnectionTypeEnum"]

squirrels/dashboards.py

@@ -1,82 +1,3 @@
-import matplotlib.figure as _figure, io as _io, abc as _abc, typing as _t
+from ._dashboards import PngDashboard, HtmlDashboard
 
-from . import _constants as _c
-
-
-class Dashboard(metaclass=_abc.ABCMeta):
-    """
-    Abstract parent class for all Dashboard classes.
-    """
-    
-    @property
-    @_abc.abstractmethod
-    def _content(self) -> bytes | str:
-        pass
-    
-    @property
-    @_abc.abstractmethod
-    def _format(self) -> str:
-        pass
-
-
-class PngDashboard(Dashboard):
-    """
-    Instantiate a Dashboard in PNG format from a matplotlib figure or bytes
-    """
-    
-    def __init__(self, content: _figure.Figure | _io.BytesIO | bytes) -> None:
-        """
-        Constructor for PngDashboard
-
-        Arguments:
-            content: The content of the dashboard as a matplotlib.figure.Figure or bytes
-        """
-        if isinstance(content, _figure.Figure):
-            buffer = _io.BytesIO()
-            content.savefig(buffer, format=_c.PNG)
-            content = buffer.getvalue()
-        
-        if isinstance(content, _io.BytesIO):
-            content = content.getvalue()
-        
-        self.__content = content
-
-    @property
-    def _content(self) -> bytes:
-        return self.__content
-    
-    @property
-    def _format(self) -> _t.Literal['png']:
-        return _c.PNG
-    
-    def _repr_png_(self):
-        return self._content
-    
-
-class HtmlDashboard(Dashboard):
-    """
-    Instantiate a Dashboard from an HTML string
-    """
-
-    def __init__(self, content: _io.StringIO | str) -> None:
-        """
-        Constructor for HtmlDashboard
-
-        Arguments:
-            content: The content of the dashboard as HTML string
-        """
-        if isinstance(content, _io.StringIO):
-            content = content.getvalue()
-        
-        self.__content = content
-
-    @property
-    def _content(self) -> str:
-        return self.__content
-    
-    @property
-    def _format(self) -> _t.Literal['html']:
-        return _c.HTML
-    
-    def _repr_html_(self):
-        return self._content
+__all__ = ["PngDashboard", "HtmlDashboard"]