squirrels 0.5.0b2__py3-none-any.whl → 0.5.0b4__py3-none-any.whl

squirrels/_project.py

@@ -1,17 +1,20 @@
 from dotenv import dotenv_values
 from uuid import uuid4
+from pathlib import Path
 import asyncio, typing as t, functools as ft, shutil, json, os
 import logging as l, matplotlib.pyplot as plt, networkx as nx, polars as pl
 import sqlglot, sqlglot.expressions
 
-from ._auth import Authenticator, BaseUser
+from ._auth import Authenticator, BaseUser, AuthProviderArgs, ProviderFunctionType
+from ._schemas import response_models as rm
 from ._model_builder import ModelBuilder
 from ._exceptions import InvalidInputError, ConfigurationError
-from . import _utils as u, _constants as c, _manifest as mf, _connection_set as cs, _api_response_models as arm
+from ._py_module import PyModule
+from . import _dashboards as d, _utils as u, _constants as c, _manifest as mf, _connection_set as cs
 from . import _seeds as s, _models as m, _model_configs as mc, _model_queries as mq, _sources as so
-from . import _parameter_sets as ps, _dashboards_io as d, dashboards as dash, dataset_result as dr
+from . import _parameter_sets as ps, _dataset_types as dr
 
-T = t.TypeVar("T", bound=dash.Dashboard)
+T = t.TypeVar("T", bound=d.Dashboard)
 M = t.TypeVar("M", bound=m.DataModel)
 
 
@@ -70,7 +73,7 @@ class SquirrelsProject:
             raise ValueError("log_format must be either 'text' or 'json'")
             
         if log_file:
-            path = u.Path(base_path, c.LOGS_FOLDER, log_file)
+            path = Path(base_path, c.LOGS_FOLDER, log_file)
             path.parent.mkdir(parents=True, exist_ok=True)
 
             handler = l.FileHandler(path)
@@ -128,12 +131,33 @@ class SquirrelsProject:
         return cs.ConnectionSetIO.load_from_file(self._logger, self._filepath, self._manifest_cfg, self._conn_args)
     
     @ft.cached_property
-    def _auth(self) -> Authenticator:
-        return Authenticator(self._logger, self._filepath, self._env_vars)
+    def _user_cls_and_provider_functions(self) -> tuple[type[BaseUser], list[ProviderFunctionType]]:
+        user_module_path = u.Path(self._filepath, c.PYCONFIGS_FOLDER, c.USER_FILE)
+        user_module = PyModule(user_module_path)
+        
+        User = user_module.get_func_or_class("User", default_attr=BaseUser)  # adds to Authenticator.providers as side effect
+        provider_functions = Authenticator.providers
+        Authenticator.providers = []
+        
+        if not issubclass(User, BaseUser):
+            raise ConfigurationError(f"User class in '{c.USER_FILE}' must inherit from BaseUser")
+        
+        return User, provider_functions
+    
+    @ft.cached_property
+    def _auth_args(self) -> AuthProviderArgs:
+        conn_args = self._conn_args
+        return AuthProviderArgs(conn_args.project_path, conn_args.proj_vars, conn_args.env_vars)
+    
+    @ft.cached_property
+    def _auth(self) -> Authenticator[BaseUser]:
+        User, provider_functions = self._user_cls_and_provider_functions
+        return Authenticator(self._logger, self._filepath, self._auth_args, provider_functions, user_cls=User)
     
     @ft.cached_property
     def _param_args(self) -> ps.ParametersArgs:
-        return ps.ParameterConfigsSetIO.get_param_args(self._conn_args)
+        conn_args = self._conn_args
+        return ps.ParametersArgs(conn_args.project_path, conn_args.proj_vars, conn_args.env_vars)
     
     @ft.cached_property
     def _param_cfg_set(self) -> ps.ParameterConfigsSet:
@@ -143,12 +167,32 @@ class SquirrelsProject:
     
     @ft.cached_property
     def _j2_env(self) -> u.EnvironmentWithMacros:
-        return u.EnvironmentWithMacros(self._logger, loader=u.j2.FileSystemLoader(self._filepath))
+        env = u.EnvironmentWithMacros(self._logger, loader=u.j2.FileSystemLoader(self._filepath))
+
+        def value_to_str(value: t.Any, attribute: str | None = None) -> str:
+            if attribute is None:
+                return str(value)
+            else:
+                return str(getattr(value, attribute))
+
+        def join(value: list[t.Any], d: str = ", ", attribute: str | None = None) -> str:
+            return d.join(map(lambda x: value_to_str(x, attribute), value))
+
+        def quote(value: t.Any, q: str = "'", attribute: str | None = None) -> str:
+            return q + value_to_str(value, attribute) + q
+        
+        def quote_and_join(value: list[t.Any], q: str = "'", d: str = ", ", attribute: str | None = None) -> str:
+            return d.join(map(lambda x: quote(x, q, attribute), value))
+        
+        env.filters["join"] = join
+        env.filters["quote"] = quote
+        env.filters["quote_and_join"] = quote_and_join
+        return env
 
     @ft.cached_property
     def _duckdb_venv_path(self) -> str:
         duckdb_filepath_setting_val = self._env_vars.get(c.SQRL_DUCKDB_VENV_DB_FILE_PATH, f"{c.TARGET_FOLDER}/{c.DUCKDB_VENV_FILE}")
-        return str(u.Path(self._filepath, duckdb_filepath_setting_val))
+        return str(Path(self._filepath, duckdb_filepath_setting_val))
     
     def close(self) -> None:
         """
@@ -236,7 +280,7 @@ class SquirrelsProject:
             for model_name in dependencies:
                 model = models_dict[model_name]
                 if isinstance(model, m.SourceModel) and not model.model_config.load_to_duckdb:
-                    raise InvalidInputError(203, f"Source model '{model_name}' cannot be queried with DuckDB")
+                    raise InvalidInputError(400, "Unqueryable source model", f"Source model '{model_name}' cannot be queried with DuckDB")
                 if isinstance(model, (m.SourceModel, m.BuildModel)):
                     substitutions[model_name] = f"venv.{model_name}"
             
@@ -255,7 +299,7 @@ class SquirrelsProject:
         dag = m.DAG(None, fake_target_model, models_dict, self._duckdb_venv_path, self._logger)
         return dag
     
-    def _draw_dag(self, dag: m.DAG, output_folder: u.Path) -> None:
+    def _draw_dag(self, dag: m.DAG, output_folder: Path) -> None:
         color_map = {
             m.ModelType.SEED: "green", m.ModelType.DBVIEW: "red", m.ModelType.FEDERATE: "skyblue",
             m.ModelType.BUILD: "purple", m.ModelType.SOURCE: "orange"
@@ -275,7 +319,7 @@ class SquirrelsProject:
         
         fig.tight_layout()
         plt.margins(x=0.1, y=0.1)
-        fig.savefig(u.Path(output_folder, "dag.png"))
+        fig.savefig(Path(output_folder, "dag.png"))
         plt.close(fig)
     
     async def _get_compiled_dag(self, *, sql_query: str | None = None, selections: dict[str, t.Any] = {}, user: BaseUser | None = None) -> m.DAG:
@@ -285,18 +329,18 @@ class SquirrelsProject:
         await dag.execute(self._param_args, self._param_cfg_set, self._context_func, user, selections, runquery=False, default_traits=default_traits)
         return dag
     
-    def _get_all_connections(self) -> list[arm.ConnectionItemModel]:
+    def _get_all_connections(self) -> list[rm.ConnectionItemModel]:
         connections = []
         for conn_name, conn_props in self._conn_set.get_connections_as_dict().items():
             if isinstance(conn_props, mf.ConnectionProperties):
                 label = conn_props.label if conn_props.label is not None else conn_name
-                connections.append(arm.ConnectionItemModel(name=conn_name, label=label))
+                connections.append(rm.ConnectionItemModel(name=conn_name, label=label))
         return connections
     
-    def _get_all_data_models(self, compiled_dag: m.DAG) -> list[arm.DataModelItem]:
+    def _get_all_data_models(self, compiled_dag: m.DAG) -> list[rm.DataModelItem]:
         return compiled_dag.get_all_data_models()
     
-    async def get_all_data_models(self) -> list[arm.DataModelItem]:
+    async def get_all_data_models(self) -> list[rm.DataModelItem]:
         """
         Get all data models in the project
 
@@ -306,26 +350,26 @@ class SquirrelsProject:
         compiled_dag = await self._get_compiled_dag()
         return self._get_all_data_models(compiled_dag)
 
-    def _get_all_data_lineage(self, compiled_dag: m.DAG) -> list[arm.LineageRelation]:
+    def _get_all_data_lineage(self, compiled_dag: m.DAG) -> list[rm.LineageRelation]:
         all_lineage = compiled_dag.get_all_model_lineage()
 
         # Add dataset nodes to the lineage
         for dataset in self._manifest_cfg.datasets.values():
-            target_dataset = arm.LineageNode(name=dataset.name, type="dataset")
-            source_model = arm.LineageNode(name=dataset.model, type="model")
-            all_lineage.append(arm.LineageRelation(type="runtime", source=source_model, target=target_dataset))
+            target_dataset = rm.LineageNode(name=dataset.name, type="dataset")
+            source_model = rm.LineageNode(name=dataset.model, type="model")
+            all_lineage.append(rm.LineageRelation(type="runtime", source=source_model, target=target_dataset))
 
         # Add dashboard nodes to the lineage
         for dashboard in self._dashboards.values():
-            target_dashboard = arm.LineageNode(name=dashboard.dashboard_name, type="dashboard")
+            target_dashboard = rm.LineageNode(name=dashboard.dashboard_name, type="dashboard")
             datasets = set(x.dataset for x in dashboard.config.depends_on)
             for dataset in datasets:
-                source_dataset = arm.LineageNode(name=dataset, type="dataset")
-                all_lineage.append(arm.LineageRelation(type="runtime", source=source_dataset, target=target_dashboard))
+                source_dataset = rm.LineageNode(name=dataset, type="dataset")
+                all_lineage.append(rm.LineageRelation(type="runtime", source=source_dataset, target=target_dashboard))
 
         return all_lineage
 
-    async def get_all_data_lineage(self) -> list[arm.LineageRelation]:
+    async def get_all_data_lineage(self) -> list[rm.LineageRelation]:
         """
         Get all data lineage in the project
 
@@ -371,28 +415,28 @@ class SquirrelsProject:
             runquery=runquery, recurse=recurse, default_traits=self._manifest_cfg.get_default_traits()
         )
         
-        output_folder = u.Path(self._filepath, c.TARGET_FOLDER, c.COMPILE_FOLDER, dataset, test_set)
+        output_folder = Path(self._filepath, c.TARGET_FOLDER, c.COMPILE_FOLDER, dataset, test_set)
         if output_folder.exists():
             shutil.rmtree(output_folder)
         output_folder.mkdir(parents=True, exist_ok=True)
         
         def write_placeholders() -> None:
-            output_filepath = u.Path(output_folder, "placeholders.json")
+            output_filepath = Path(output_folder, "placeholders.json")
             with open(output_filepath, 'w') as f:
                 json.dump(dag.placeholders, f, indent=4)
         
         def write_model_outputs(model: m.DataModel) -> None:
             assert isinstance(model, m.QueryModel)
             subfolder = c.DBVIEWS_FOLDER if model.model_type == m.ModelType.DBVIEW else c.FEDERATES_FOLDER
-            subpath = u.Path(output_folder, subfolder)
+            subpath = Path(output_folder, subfolder)
             subpath.mkdir(parents=True, exist_ok=True)
             if isinstance(model.compiled_query, mq.SqlModelQuery):
-                output_filepath = u.Path(subpath, model.name+'.sql')
+                output_filepath = Path(subpath, model.name+'.sql')
                 query = model.compiled_query.query
                 with open(output_filepath, 'w') as f:
                     f.write(query)
             if runquery and isinstance(model.result, pl.LazyFrame):
-                output_filepath = u.Path(subpath, model.name+'.csv')
+                output_filepath = Path(subpath, model.name+'.csv')
                 model.result.collect().write_csv(output_filepath)
 
         write_placeholders()
@@ -455,7 +499,7 @@ class SquirrelsProject:
 
     def _permission_error(self, user: BaseUser | None, data_type: str, data_name: str, scope: str) -> InvalidInputError:
         username = "" if user is None else f" '{user.username}'"
-        return InvalidInputError(25, f"User{username} does not have permission to access {scope} {data_type}: {data_name}")
+        return InvalidInputError(403, f"Unauthorized access to {data_type}", f"User{username} does not have permission to access {scope} {data_type}: {data_name}")
     
     def seed(self, name: str) -> pl.LazyFrame:
         """
@@ -520,7 +564,7 @@ class SquirrelsProject:
         )
     
     async def dashboard(
-        self, name: str, *, selections: dict[str, t.Any] = {}, user: BaseUser | None = None, dashboard_type: t.Type[T] = dash.Dashboard
+        self, name: str, *, selections: dict[str, t.Any] = {}, user: BaseUser | None = None, dashboard_type: t.Type[T] = d.PngDashboard
     ) -> T:
         """
         Async method to retrieve a dashboard given parameter selections.

squirrels/_py_module.py

@@ -43,11 +43,12 @@ class PyModule:
         return func_or_class
 
 
-def run_pyconfig_main(base_path: str, filename: str, kwargs: dict[str, Any] = {}) -> None:
+def run_pyconfig_main(base_path: str, filename: str, kwargs: dict[str, Any] = {}) -> Any | None:
     """
     Given a python file in the 'pyconfigs' folder, run its main function
     
     Arguments:
+        base_path: The base path of the project
         filename: The name of the file to run main function
         kwargs: Dictionary of the main function arguments
     """
@@ -56,6 +57,6 @@ def run_pyconfig_main(base_path: str, filename: str, kwargs: dict[str, Any] = {}
     main_function = module.get_func_or_class(c.MAIN_FUNC, is_required=False)
     if main_function:
         try:
-            main_function(**kwargs)
+            return main_function(**kwargs)
         except Exception as e:
             raise FileExecutionError(f'Failed to run python file "{filepath}"', e) from e

squirrels/_schemas/auth_models.py

@@ -0,0 +1,144 @@
+from typing import Callable, Any
+from datetime import datetime
+from pydantic import BaseModel, ConfigDict, Field, field_serializer
+
+
+class BaseUser(BaseModel):
+    model_config = ConfigDict(from_attributes=True)
+    username: str
+    is_admin: bool = False
+    
+    @classmethod
+    def dropped_columns(cls):
+        return []
+    
+    def __hash__(self):
+        return hash(self.username)
+
+
+class ApiKey(BaseModel):
+    model_config = ConfigDict(from_attributes=True)
+    id: str
+    title: str
+    username: str
+    created_at: datetime
+    expires_at: datetime
+    
+    @field_serializer('created_at', 'expires_at')
+    def serialize_datetime(self, dt: datetime) -> str:
+        return dt.strftime("%Y-%m-%dT%H:%M:%S.%fZ")
+
+
+class UserField(BaseModel):
+    name: str
+    type: str
+    nullable: bool
+    enum: list[str] | None
+    default: Any | None
+
+
+class ProviderConfigs(BaseModel):
+    client_id: str
+    client_secret: str
+    server_metadata_url: str
+    client_kwargs: dict = Field(default_factory=dict)
+    get_user: Callable[[dict], BaseUser]
+
+
+class AuthProvider(BaseModel):
+    name: str
+    label: str
+    icon: str
+    provider_configs: ProviderConfigs
+
+
+# OAuth 2.1 Models
+
+class OAuthClientModel(BaseModel):
+    """OAuth client details"""
+    model_config = ConfigDict(from_attributes=True)
+    client_id: str
+    client_name: str
+    redirect_uris: list[str]
+    scope: str
+    grant_types: list[str]
+    response_types: list[str]
+    created_at: datetime
+    is_active: bool
+    
+    @field_serializer('created_at')
+    def serialize_datetime(self, dt: datetime) -> str:
+        return dt.strftime("%Y-%m-%dT%H:%M:%S.%fZ")
+
+
+class ClientRegistrationRequest(BaseModel):
+    """Request model for OAuth client registration"""
+    client_name: str = Field(description="Human-readable name for the OAuth client")
+    redirect_uris: list[str] = Field(description="List of allowed redirect URIs for the client")
+    scope: str = Field(default="read", description="Default scope for the client")
+    grant_types: list[str] = Field(default=["authorization_code", "refresh_token"], description="Allowed grant types")
+    response_types: list[str] = Field(default=["code"], description="Allowed response types")
+
+
+class ClientUpdateRequest(BaseModel):
+    """Request model for OAuth client update"""
+    client_name: str | None = Field(default=None, description="Human-readable name for the OAuth client")
+    redirect_uris: list[str] | None = Field(default=None, description="List of allowed redirect URIs for the client")
+    scope: str | None = Field(default=None, description="Default scope for the client")
+    grant_types: list[str] | None = Field(default=None, description="Allowed grant types")
+    response_types: list[str] | None = Field(default=None, description="Allowed response types")
+    is_active: bool | None = Field(default=None, description="Whether the client is active")
+
+
+class ClientDetailsResponse(BaseModel):
+    """Response model for OAuth client details (without client_secret)"""
+    client_id: str = Field(description="Client ID")
+    client_name: str = Field(description="Client name")
+    redirect_uris: list[str] = Field(description="Registered redirect URIs")
+    scope: str = Field(description="Default scope")
+    grant_types: list[str] = Field(description="Allowed grant types")
+    response_types: list[str] = Field(description="Allowed response types")
+    created_at: datetime = Field(description="Registration timestamp")
+    is_active: bool = Field(description="Whether the client is active")
+    
+    @field_serializer('created_at')
+    def serialize_datetime(self, dt: datetime) -> str:
+        return dt.strftime("%Y-%m-%dT%H:%M:%S.%fZ")
+
+
+class ClientUpdateResponse(ClientDetailsResponse):
+    """Response model for OAuth client update"""
+    registration_access_token: str | None = Field(default=None, description="Token for managing this client registration (store securely)")
+
+
+class ClientRegistrationResponse(ClientUpdateResponse):
+    """Response model for OAuth client registration"""
+    client_secret: str = Field(description="Generated client secret (store securely)")
+    registration_client_uri: str | None = Field(default=None, description="URI for managing this client registration")
+
+
+class TokenResponse(BaseModel):
+    access_token: str
+    token_type: str = "bearer"
+    expires_in: int
+    refresh_token: str | None = None
+
+
+class OAuthServerMetadata(BaseModel):
+    """OAuth 2.1 Authorization Server Metadata (RFC 8414)"""
+    issuer: str = Field(description="Authorization server's issuer identifier URL")
+    authorization_endpoint: str = Field(description="URL of the authorization endpoint")
+    token_endpoint: str = Field(description="URL of the token endpoint")
+    revocation_endpoint: str = Field(description="URL of the token revocation endpoint")
+    registration_endpoint: str = Field(description="URL of the client registration endpoint")
+    scopes_supported: list[str] = Field(description="List of OAuth 2.1 scope values supported")
+    response_types_supported: list[str] = Field(description="List of OAuth 2.1 response_type values supported")
+    grant_types_supported: list[str] = Field(description="List of OAuth 2.1 grant type values supported")
+    token_endpoint_auth_methods_supported: list[str] = Field(
+        default=["client_secret_basic", "client_secret_post"], 
+        description="List of client authentication methods supported by the token endpoint"
+    )
+    code_challenge_methods_supported: list[str] = Field(
+        default=["S256"], 
+        description="List of PKCE code challenge methods supported"
+    )

squirrels/_schemas/query_param_models.py

@@ -0,0 +1,67 @@
+"""
+Query model generation utilities for API routes
+"""
+from typing import Annotated
+from dataclasses import make_dataclass
+from fastapi import Depends
+from pydantic import create_model
+
+from .._parameter_configs import APIParamFieldInfo
+
+
+def _get_query_models_helper(widget_parameters: list[str] | None, predefined_params: list[APIParamFieldInfo], param_fields: dict):
+    """Helper function to generate query models"""
+    if widget_parameters is None:
+        widget_parameters = list(param_fields.keys())
+    
+    QueryModelForGetRaw = make_dataclass("QueryParams", [
+        param_fields[param].as_query_info() for param in widget_parameters
+    ] + [param.as_query_info() for param in predefined_params])
+    QueryModelForGet = Annotated[QueryModelForGetRaw, Depends()]
+
+    field_definitions = {param: param_fields[param].as_body_info() for param in widget_parameters}
+    for param in predefined_params:
+        field_definitions[param.name] = param.as_body_info()
+    QueryModelForPost = create_model("RequestBodyParams", **field_definitions) # type: ignore
+    return QueryModelForGet, QueryModelForPost
+
+
+def get_query_models_for_parameters(widget_parameters: list[str] | None, param_fields: dict):
+    """Generate query models for parameter endpoints"""
+    predefined_params = [
+        APIParamFieldInfo("x_verify_params", bool, default=False, description="If true, the query parameters are verified to be valid for the dataset"),
+        APIParamFieldInfo("x_parent_param", str, description="The parameter name used for parameter updates. If not provided, then all parameters are retrieved"),
+    ]
+    return _get_query_models_helper(widget_parameters, predefined_params, param_fields)
+
+
+def get_query_models_for_dataset(widget_parameters: list[str] | None, param_fields: dict):
+    """Generate query models for dataset endpoints"""
+    predefined_params = [
+        APIParamFieldInfo("x_verify_params", bool, default=False, description="If true, the query parameters are verified to be valid for the dataset"),
+        APIParamFieldInfo("x_orientation", str, default="records", description="The orientation of the data to return, one of: 'records', 'rows', or 'columns'"),
+        APIParamFieldInfo("x_select", list[str], examples=[[]], description="The columns to select from the dataset. All are returned if not specified"), 
+        APIParamFieldInfo("x_offset", int, default=0, description="The number of rows to skip before returning data (applied after data caching)"),
+        APIParamFieldInfo("x_limit", int, default=1000, description="The maximum number of rows to return (applied after data caching and offset)"),
+    ]
+    return _get_query_models_helper(widget_parameters, predefined_params, param_fields)
+
+
+def get_query_models_for_dashboard(widget_parameters: list[str] | None, param_fields: dict):
+    """Generate query models for dashboard endpoints"""
+    predefined_params = [
+        APIParamFieldInfo("x_verify_params", bool, default=False, description="If true, the query parameters are verified to be valid for the dashboard"),
+    ]
+    return _get_query_models_helper(widget_parameters, predefined_params, param_fields)
+
+
+def get_query_models_for_querying_models(param_fields: dict):
+    """Generate query models for querying data models"""
+    predefined_params = [
+        APIParamFieldInfo("x_verify_params", bool, default=False, description="If true, the query parameters are verified to be valid"),
+        APIParamFieldInfo("x_orientation", str, default="records", description="The orientation of the data to return, one of: 'records', 'rows', or 'columns'"),
+        APIParamFieldInfo("x_offset", int, default=0, description="The number of rows to skip before returning data (applied after data caching)"),
+        APIParamFieldInfo("x_limit", int, default=1000, description="The maximum number of rows to return (applied after data caching and offset)"),
+        APIParamFieldInfo("x_sql_query", str, description="The SQL query to execute on the data models"),
+    ]
+    return _get_query_models_helper(None, predefined_params, param_fields) 

squirrels/{_api_response_models.py → _schemas/response_models.py}

@@ -1,16 +1,20 @@
 from typing import Annotated, Literal
 from pydantic import BaseModel, Field
-from datetime import datetime, date
+from datetime import date
 
-from . import _model_configs as mc, _sources as s
+from .. import _model_configs as mc, _sources as s
 
 
-class LoginReponse(BaseModel):
-    access_token: Annotated[str, Field(examples=["encoded_jwt_token"], description="An encoded JSON web token to use subsequent API requests")]
-    token_type: Annotated[str, Field(examples=["bearer"], description='Always "bearer" for Bearer token')]
-    username: Annotated[str, Field(examples=["johndoe"], description='The username authenticated with from the form data')]
-    is_admin: Annotated[bool, Field(examples=[False], description="A boolean for whether the user is an admin")]
-    expiry_time: Annotated[datetime, Field(examples=["2023-08-01T12:00:00.000000Z"], description="The expiry time of the access token in yyyy-MM-dd'T'hh:mm:ss.SSSSSS'Z' format")]
+## Simple Auth Response Models
+
+class ApiKeyResponse(BaseModel):
+    api_key: Annotated[str, Field(examples=["sqrl-12345678"], description="The API key to use subsequent API requests")]
+
+class ProviderResponse(BaseModel):
+    name: Annotated[str, Field(examples=["my_provider"], description="The name of the provider")]
+    label: Annotated[str, Field(examples=["My Provider"], description="The human-friendly display name for the provider")]
+    icon: Annotated[str, Field(examples=["https://example.com/my_provider_icon.png"], description="The URL of the provider's icon")]
+    login_url: Annotated[str, Field(examples=["https://example.com/my_provider_login"], description="The URL to redirect to for provider login")]
 
 
 ## Parameters Response Models

squirrels/_utils.py

@@ -2,10 +2,9 @@ from typing import Sequence, Optional, Union, TypeVar, Callable, Any, Iterable
 from datetime import datetime
 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
+import sqlglot, sqlglot.expressions, asyncio, hashlib, inspect, base64
 
 from . import _constants as c
 from ._exceptions import ConfigurationError
@@ -290,7 +289,8 @@ def load_yaml_config(filepath: FilePath) -> dict:
 
 
 def run_duckdb_stmt(
-    logger: Logger, duckdb_conn: duckdb.DuckDBPyConnection, stmt: str, *, params: dict[str, Any] | None = None, redacted_values: list[str] = []
+    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
@@ -306,7 +306,8 @@ def run_duckdb_stmt(
     for value in redacted_values:
         redacted_stmt = redacted_stmt.replace(value, "[REDACTED]")
     
-    logger.info(f"Running statement: {redacted_stmt}", extra={"data": {"params": params}})
+    for_model_name = f" for model '{model_name}'" if model_name is not None else ""
+    logger.info(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:
@@ -357,3 +358,36 @@ async def asyncio_gather(coroutines: list):
         # 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

squirrels/arguments.py

@@ -0,0 +1,2 @@
+from ._arguments.init_time_args import ConnectionsArgs, AuthProviderArgs, ParametersArgs, BuildModelArgs
+from ._arguments.run_time_args import ContextArgs, ModelArgs, DashboardArgs

squirrels/auth.py

@@ -0,0 +1 @@
+from ._auth import BaseUser, ProviderConfigs, provider

squirrels/connections.py

@@ -0,0 +1 @@
+from ._manifest import ConnectionProperties, ConnectionTypeEnum