squirrels 0.2.2__py3-none-any.whl → 0.3.0__py3-none-any.whl

squirrels/_utils.py

@@ -1,7 +1,6 @@
-from typing import Sequence, Optional, Union, Any, TypeVar, Callable
+from typing import Sequence, Optional, Union, TypeVar, Callable
 from pathlib import Path
-from pandas.api import types as pd_types
-import json, jinja2 as j2, pandas as pd
+import json, sqlite3, jinja2 as j2, pandas as pd
 
 from . import _constants as c
 
@@ -24,7 +23,7 @@ class ConfigurationError(Exception):
 
 class FileExecutionError(ConfigurationError):
     def __init__(self, message: str, error: Exception, *args) -> None:
-        new_message = message + f"\n... Produced error message `{error}` (see above for more details)"
+        new_message = message + f"\n... Produced error message `{error}` (scroll up for more details on handled exception)"
         super().__init__(new_message, *args)
 
 
@@ -45,7 +44,7 @@ def join_paths(*paths: FilePath) -> Path:
 
 _j2_env = j2.Environment(loader=j2.FileSystemLoader('.'))
 
-def render_string(raw_str: str, kwargs: dict) -> str:
+def render_string(raw_str: str, **kwargs: dict) -> str:
     """
     Given a template string, render it with the given keyword arguments
 
@@ -112,33 +111,6 @@ def normalize_name_for_api(name: str) -> str:
     return name.replace('_', '-')
 
 
-def df_to_json0(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 {"schema": out_schema, "data": in_df_json["data"]}
-
-
 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
@@ -183,6 +155,42 @@ def process_if_not_none(input_val: Optional[X], processor: Callable[[X], Y]) ->
     return processor(input_val)
 
 
-def use_duckdb():
+def use_duckdb() -> bool:
+    """
+    Determines whether to use DuckDB instead of SQLite for embedded database
+
+    Returns:
+        A boolean
+    """
     from ._manifest import ManifestIO
     return (ManifestIO.obj.settings.get(c.IN_MEMORY_DB_SETTING, c.SQLITE) == c.DUCKDB)
+
+
+def run_sql_on_dataframes(sql_query: str, dataframes: dict[str, pd.DataFrame]) -> pd.DataFrame:
+    """
+    Runs a SQL query against a collection of dataframes
+
+    Parameters:
+        sql_query: The SQL query to run
+        dataframes: A dictionary of table names to their pandas Dataframe
+    
+    Returns:
+        The result as a pandas Dataframe from running the query
+    """
+    do_use_duckdb = use_duckdb()
+    if do_use_duckdb:
+        import duckdb
+        conn = duckdb.connect()
+    else:
+        conn = sqlite3.connect(":memory:")
+    
+    try:
+        for name, df in dataframes.items():
+            if do_use_duckdb:
+                conn.execute(f"CREATE TABLE {name} AS FROM df")
+            else:
+                df.to_sql(name, conn, index=False)
+        
+        return conn.execute(sql_query).df() if do_use_duckdb else pd.read_sql(sql_query, conn)
+    finally:
+        conn.close()

squirrels/arguments/run_time_args.py

@@ -1,11 +1,11 @@
-from typing import Callable, Any
+from typing import Union, Callable, Optional, Any
 from dataclasses import dataclass
 from sqlalchemy import Engine
-import pandas as pd, sqlite3
+import pandas as pd
 
 from .init_time_args import ConnectionsArgs, ParametersArgs
 from ..user_base import User
-from ..parameters import Parameter
+from ..parameters import Parameter, _TextValue
 from .._connection_set import ConnectionSetIO
 from .. import _utils as u
 
@@ -22,6 +22,32 @@ class ContextArgs(ParametersArgs):
     user: User
     prms: dict[str, Parameter]
     traits: dict[str, Any]
+    _placeholders: dict[str, Any]
+
+    def set_placeholder(self, placeholder: str, value: Union[_TextValue, Any]) -> None:
+        """
+        Method to set a placeholder value.
+
+        Parameters:
+            placeholder: A string for the name of the placeholder
+            value: The value of the placeholder. Can be of any type
+        """
+        if isinstance(value, _TextValue):
+            value = value._value_do_not_touch
+        self._placeholders[placeholder] = value
+    
+    def param_exists(self, param_name: str) -> bool:
+        """
+        Method to check whether a given parameter exists and is enabled (i.e., not hidden based on other parameter selections) for the current 
+        dataset at runtime.
+
+        Parameters:
+            param_name: A string for the name of the parameter
+        
+        Returns:
+            A boolean for whether the parameter exists
+        """
+        return (param_name in self.prms and self.prms[param_name].is_enabled())
 
 
 @dataclass
@@ -32,12 +58,43 @@ class ModelDepsArgs(ContextArgs):
 @dataclass
 class ModelArgs(ModelDepsArgs):
     connection_name: str
-    connections: dict[str, Engine]
+    _connections: dict[str, Engine]
+    _dependencies: set[str]
     _ref: Callable[[str], pd.DataFrame]
-    dependencies: set[str]
 
-    def __post_init__(self):
-        self.ref = self._ref
+    @property
+    def connections(self) -> dict[str, Engine]:
+        return self._connections.copy()
+
+    @property
+    def dependencies(self) -> set[str]:
+        return self._dependencies.copy()
+    
+    def is_placeholder(self, placeholder: str) -> bool:
+        """
+        Checks whether a name is a valid placeholder
+
+        Parameters:
+            placeholder: A string for the name of the placeholder
+        
+        Returns:
+            A boolean for whether name exists
+        """
+        return placeholder in self._placeholders
+    
+    def get_placeholder_value(self, placeholder: str) -> Optional[Any]:
+        """
+        Gets the value of a placeholder.
+
+        USE WITH CAUTION. Do not use the return value directly in a SQL query since that could be prone to SQL injection
+
+        Parameters:
+            placeholder: A string for the name of the placeholder
+        
+        Returns:
+            An type for the value of the placeholder
+        """
+        return self._placeholders.get(placeholder)
     
     def ref(self, model: str) -> pd.DataFrame:
         """
@@ -52,49 +109,34 @@ class ModelArgs(ModelDepsArgs):
         Returns:
             A pandas DataFrame
         """
+        return self._ref(model)
 
-    def run_external_sql(self, sql: str, *, connection_name: str = None, **kwargs) -> pd.DataFrame:
+    def run_external_sql(self, sql_query: str, *, connection_name: str = None, **kwargs) -> pd.DataFrame:
         """
-        Runs a SQL query against an external database, with option to specify the connection name
+        Runs a SQL query against an external database, with option to specify the connection name. Placeholder values are provided automatically
 
         Parameters:
-            sql: The SQL query
+            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 pandas DataFrame
         """
         connection_name = self.connection_name if connection_name is None else connection_name
-        return ConnectionSetIO.obj.run_sql_query_from_conn_name(sql, connection_name)
+        return ConnectionSetIO.obj.run_sql_query_from_conn_name(sql_query, connection_name, self._placeholders)
 
-    def run_sql_on_dataframes(self, query: str, *, dataframes: dict[str, pd.DataFrame] = None, **kwargs) -> pd.DataFrame:
+    def run_sql_on_dataframes(self, sql_query: str, *, dataframes: dict[str, pd.DataFrame] = None, **kwargs) -> pd.DataFrame:
         """
-        Uses a dictionary of dataframes to execute a SQL query in an in-memory sqlite database
+        Uses a dictionary of dataframes to execute a SQL query in an embedded in-memory database (sqlite or duckdb based on setting)
 
         Parameters:
-            query: The SQL query to run using sqlite
-            dataframes: A dictionary of table names to their pandas Dataframe
+            sql_query: The SQL query to run
+            dataframes: A dictionary of table names to their pandas Dataframe. If None, uses results of dependent models
         
         Returns:
             The result as a pandas Dataframe from running the query
         """
         if dataframes is None:
-            dataframes = {x: self.ref(x) for x in self.dependencies}
-
-        use_duckdb = u.use_duckdb()
-        if use_duckdb:
-            import duckdb
-            conn = duckdb.connect()
-        else:
-            conn = sqlite3.connect(":memory:")
-        
-        try:
-            for name, df in dataframes.items():
-                if use_duckdb:
-                    conn.execute(f"CREATE TABLE {name} AS FROM df")
-                else:
-                    df.to_sql(name, conn, index=False)
-            
-            return conn.execute(query).df() if use_duckdb else pd.read_sql(query, conn)
-        finally:
-            conn.close()
+            dataframes = {x: self.ref(x) for x in self._dependencies}
+
+        return u.run_sql_on_dataframes(sql_query, dataframes)