squirrels 0.4.1__py3-none-any.whl → 0.5.0rc0__py3-none-any.whl

squirrels/_sources.py

@@ -0,0 +1,106 @@
+from typing import Any
+from pydantic import BaseModel, Field, model_validator
+import time, sqlglot
+
+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 value is set")
+    selective_overwrite_value: Any = Field(default=None)
+
+
+class Source(mc.ConnectionInterface, mc.ModelConfig):
+    table: str | None = Field(default=None)
+    load_to_duckdb: bool = Field(default=False, description="Whether to load the data to DuckDB")
+    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])
+        primary_key_clause = f", PRIMARY KEY ({', '.join(self.primary_key)})" if self.primary_key else ""
+        return f"{cols_clause}{primary_key_clause}"
+    
+    def get_cols_for_insert_stmt(self) -> str:
+        return ", ".join([col.name for col in self.columns])
+    
+    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_insert(self, dialect: str, conn_name: str, table_name: str, max_value_of_increasing_col: Any | None, *, full_refresh: bool = True) -> str:
+        select_cols = self.get_cols_for_insert_stmt()
+        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}"
+        pushdown_query = f"SELECT {select_cols} FROM {table_name} WHERE {where_cond}"
+        
+        if dialect in ['postgres', 'mysql']:
+            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}"
+    
+    def get_insert_replace_clause(self) -> str:
+        return "" if len(self.primary_key) == 0 else "OR REPLACE"
+
+
+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)
+        sources_data = u.load_yaml_config(sources_path) if sources_path.exists() else {}
+        
+        sources = Sources(**sources_data).finalize_null_fields(env_vars)
+        
+        logger.log_activity_time("loading sources", start)
+        return sources

squirrels/_utils.py

@@ -1,35 +1,35 @@
-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, Any, Iterable
 from datetime import datetime
-import os, time, logging, json, sqlite3, pandas as pd
+from pathlib import Path
+from functools import lru_cache
+from pydantic import BaseModel
+import os, time, logging, json, duckdb, polars as pl, yaml
 import jinja2 as j2, jinja2.nodes as j2_nodes
+import sqlglot, sqlglot.expressions, asyncio
 
 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"],
+    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 +40,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):
@@ -115,7 +115,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
@@ -180,7 +179,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 +196,164 @@ 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:
+@lru_cache(maxsize=1)
+def _read_duckdb_init_sql() -> tuple[str, Path | None]:
+    """
+    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())
+        
+        init_sql = "\n".join(init_contents).strip()
+        target_init_path = None
+        if init_sql:
+            target_init_path = Path(c.TARGET_FOLDER, c.DUCKDB_INIT_FILE)
+            target_init_path.parent.mkdir(parents=True, exist_ok=True)
+            target_init_path.write_text(init_sql)
+        
+        return init_sql, target_init_path
+    except Exception as e:
+        raise ConfigurationError(f"Failed to read {c.DUCKDB_INIT_FILE}: {str(e)}") from e
+
+def create_duckdb_connection(filepath: str | Path = ":memory:", *, read_only: bool = False) -> 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.
+        read_only: Whether to open the database in read-only mode. Defaults to False.
+    
+    Returns:
+        A DuckDB connection (which must be closed after use)
+    """
+    conn = duckdb.connect(filepath, read_only=read_only)
+    
+    try:
+        init_sql, _ = _read_duckdb_init_sql()
+        if init_sql:
+            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:
+        A dictionary representation of the YAML file
+    """
+    try:
+        with open(filepath, 'r') as f:
+            return yaml.safe_load(f)
+    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, 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]")
+    
+    logger.info(f"Running statement: {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 response of a Squirrels dataset result API
+        The set of dependent tables
     """
-    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)
+    # Parse the SQL query and extract all table references
+    parsed = sqlglot.parse_one(sql_query)
+    dependencies = set()
     
-    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
+    # 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

squirrels/_version.py

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

squirrels/arguments/init_time_args.py

@@ -1,40 +1,103 @@
-from typing import Callable, Any
+from typing import Any, Iterable, Callable
 from dataclasses import dataclass
+import polars as pl
 
+from .. import _utils as u
 
 @dataclass
-class BaseArguments:
+class ConnectionsArgs:
+    project_path: str
     _proj_vars: dict[str, Any]
-    _env_vars: dict[str, Any]
+    _env_vars: dict[str, str]
 
     @property
     def proj_vars(self) -> dict[str, Any]:
         return self._proj_vars.copy()
     
     @property
-    def env_vars(self) -> dict[str, Any]:
+    def env_vars(self) -> dict[str, str]:
         return self._env_vars.copy()
 
 
 @dataclass
-class ConnectionsArgs(BaseArguments):
-    _get_credential: Callable[[str | None], tuple[str, str]]
+class ParametersArgs(ConnectionsArgs):
+    pass
+
+
+@dataclass
+class _WithConnectionDictArgs(ConnectionsArgs):
+    _connections: dict[str, Any]
+
+    @property
+    def connections(self) -> dict[str, Any]:
+        """
+        A dictionary of connection keys to SQLAlchemy Engines for database connections. 
+        
+        Can also be used to store other in-memory objects in advance such as ML models.
+        """
+        return self._connections.copy()
+
+
+class BuildModelArgs(_WithConnectionDictArgs):
+
+    def __init__(
+        self, conn_args: ConnectionsArgs, _connections: dict[str, Any], 
+        dependencies: Iterable[str], 
+        ref: Callable[[str], pl.LazyFrame], 
+        run_external_sql: Callable[[str, str], pl.DataFrame]
+    ):
+        super().__init__(conn_args.project_path, conn_args.proj_vars, conn_args.env_vars, _connections)
+        self._dependencies = dependencies
+        self._ref = ref
+        self._run_external_sql = run_external_sql
 
-    def get_credential(self, key: str | None) -> tuple[str, str]:
+    @property
+    def dependencies(self) -> set[str]:
+        """
+        The set of dependent data model names
+        """
+        return set(self._dependencies)
+    
+    def ref(self, model: str) -> pl.LazyFrame:
         """
-        Return (username, password) tuple configured for credentials key in env.yaml
+        Returns the result (as polars DataFrame) of a dependent model (predefined in "dependencies" function)
 
-        If key is None, returns tuple of empty strings ("", "")
+        Note: This is different behaviour than the "ref" function for SQL models, which figures out the dependent models for you, 
+        and returns a string for the table/view name instead of a polars DataFrame.
 
         Arguments:
-            key: The credentials key
+            model: The model name
         
         Returns:
-            A tuple of 2 strings
+            A polars DataFrame
         """
-        return self._get_credential(key)
+        return self._ref(model)
 
+    def run_external_sql(self, connection_name: str, sql_query: str, **kwargs) -> pl.DataFrame:
+        """
+        Runs a SQL query against an external database, with option to specify the connection name. Placeholder values are provided automatically
 
-@dataclass
-class ParametersArgs(BaseArguments):
-    pass
+        Arguments:
+            sql_query: The SQL query. Can be parameterized with placeholders
+            connection_name: The connection name for the database. If None, uses the one configured for the model
+        
+        Returns:
+            The query result as a polars DataFrame
+        """
+        return self._run_external_sql(sql_query, connection_name)
+
+    def run_sql_on_dataframes(self, sql_query: str, *, dataframes: dict[str, pl.LazyFrame] | None = None, **kwargs) -> pl.DataFrame:
+        """
+        Uses a dictionary of dataframes to execute a SQL query in an embedded in-memory database (sqlite or duckdb based on setting)
+
+        Arguments:
+            sql_query: The SQL query to run
+            dataframes: A dictionary of table names to their polars LazyFrame. If None, uses results of dependent models
+        
+        Returns:
+            The result as a polars LazyFrame from running the query
+        """
+        if dataframes is None:
+            dataframes = {x: self.ref(x) for x in self._dependencies}
+
+        return u.run_sql_on_dataframes(sql_query, dataframes)