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

squirrels/_arguments/init_time_args.py

@@ -0,0 +1,108 @@
+from typing import Any, Iterable, Callable
+from dataclasses import dataclass
+import polars as pl
+
+from .. import _utils as u
+
+@dataclass
+class ConnectionsArgs:
+    project_path: str
+    _proj_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, str]:
+        return self._env_vars.copy()
+
+
+@dataclass
+class AuthProviderArgs(ConnectionsArgs):
+    pass
+
+
+@dataclass
+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
+
+    @property
+    def dependencies(self) -> set[str]:
+        """
+        The set of dependent data model names
+        """
+        return set(self._dependencies)
+    
+    def ref(self, model: str) -> pl.LazyFrame:
+        """
+        Returns the result (as polars DataFrame) of a dependent model (predefined in "dependencies" function)
+
+        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:
+            model: The model name
+        
+        Returns:
+            A polars DataFrame
+        """
+        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
+
+        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 DuckDB database
+
+        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 DataFrame 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)

squirrels/_arguments/run_time_args.py

@@ -0,0 +1,147 @@
+from typing import Callable, Any, Coroutine
+import polars as pl
+
+from .init_time_args import ParametersArgs, BuildModelArgs
+from .._schemas.auth_models import AbstractUser
+from .._parameters import Parameter, TextValue
+
+
+class ContextArgs(ParametersArgs):
+
+    def __init__(
+        self, param_args: ParametersArgs, 
+        user: AbstractUser, 
+        prms: dict[str, Parameter],
+        configurables: dict[str, str]
+    ):
+        super().__init__(param_args.project_path, param_args.proj_vars, param_args.env_vars)
+        self.user = user
+        self._prms = prms
+        self._configurables = configurables
+        self._placeholders = {}
+
+    @property
+    def prms(self) -> dict[str, Parameter]:
+        """
+        A dictionary of parameter names to parameter
+        """
+        return self._prms.copy()
+    
+    @property
+    def configurables(self) -> dict[str, str]:
+        """
+        A dictionary of configurable name to value (set by application)
+        """
+        return self._configurables.copy()
+
+    @property
+    def _placeholders_copy(self) -> dict[str, Any]:
+        """
+        A dictionary of placeholder name to placeholder value
+        """
+        return self._placeholders.copy()
+
+    def set_placeholder(self, placeholder: str, value: TextValue | Any) -> str:
+        """
+        Method to set a placeholder value.
+
+        Arguments:
+            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
+        return ""
+    
+    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.
+
+        Arguments:
+            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())
+
+
+class ModelArgs(BuildModelArgs, ContextArgs):
+
+    def __init__(
+        self, ctx_args: ContextArgs, build_model_args: BuildModelArgs, 
+        ctx: dict[str, Any]
+    ):
+        super(ContextArgs, self).__init__(ctx_args.project_path, ctx_args.proj_vars, ctx_args.env_vars)
+        self._project_path = ctx_args.project_path
+        self._proj_vars = ctx_args.proj_vars
+        self._env_vars = ctx_args.env_vars
+        self.user = ctx_args.user
+        self._prms = ctx_args.prms
+        self._configurables = ctx_args.configurables
+        self._placeholders = ctx_args._placeholders_copy
+        self._connections = build_model_args.connections
+        self._dependencies = build_model_args.dependencies
+        self._ref = build_model_args.ref
+        self._run_external_sql = build_model_args.run_external_sql
+        self._ctx = ctx
+
+    @property
+    def ctx(self) -> dict[str, Any]:
+        """
+        Dictionary of context variables
+        """
+        return self._ctx.copy()
+    
+    def is_placeholder(self, placeholder: str) -> bool:
+        """
+        Checks whether a name is a valid placeholder
+
+        Arguments:
+            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) -> Any | None:
+        """
+        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
+
+        Arguments:
+            placeholder: A string for the name of the placeholder
+        
+        Returns:
+            An type for the value of the placeholder
+        """
+        return self._placeholders.get(placeholder)
+
+
+class DashboardArgs(ParametersArgs):
+
+    def __init__(
+        self, param_args: ParametersArgs, 
+        get_dataset: Callable[[str, dict[str, Any]], Coroutine[Any, Any, pl.DataFrame]]
+    ):
+        super().__init__(param_args.project_path, param_args.proj_vars, param_args.env_vars)
+        self._get_dataset = get_dataset
+
+    async def dataset(self, name: str, *, fixed_parameters: dict[str, Any] = {}) -> pl.DataFrame:
+        """
+        Get dataset as DataFrame given dataset name. Can use this to access protected/private datasets regardless of user authenticated to the dashboard.
+
+        The parameters used for the dataset include the parameter selections coming from the REST API and the fixed_parameters argument. The fixed_parameters takes precedence.
+
+        Arguments:
+            name: A string for the dataset name
+            fixed_parameters: Parameters to set for this dataset (in addition to the ones set through real-time selections)
+        
+        Returns:
+            A DataFrame for the result of the dataset
+        """
+        return await self._get_dataset(name, fixed_parameters)