squirrels 0.5.0b3__py3-none-any.whl → 0.6.0.post0__py3-none-any.whl

squirrels/_env_vars.py

@@ -0,0 +1,209 @@
+from typing import Any, Literal, Optional
+from typing_extensions import Self
+from pathlib import Path
+from pydantic import BaseModel, Field, field_validator, model_validator, ConfigDict
+import json
+
+from . import _constants as c
+
+
+class SquirrelsEnvVars(BaseModel):
+    """
+    Pydantic model for managing and validating Squirrels environment variables.
+    These variables are typically loaded from .env files or the system environment.
+    """
+    model_config = ConfigDict(serialize_by_alias=True)
+    project_path: str
+    
+    # Security
+    secret_key: Optional[str] = Field(
+        None, alias=c.SQRL_SECRET_KEY, 
+        description="Secret key for JWT encoding/decoding and other security operations"
+    )
+    secret_admin_password: Optional[str] = Field(
+        None, alias=c.SQRL_SECRET_ADMIN_PASSWORD, 
+        description="Password for the admin user"
+    )
+    
+    # Auth
+    auth_db_file_path: str = Field(
+        "{project_path}/target/auth.sqlite", alias=c.SQRL_AUTH_DB_FILE_PATH, 
+        description="Path to the SQLite authentication database"
+    )
+    auth_token_expire_minutes: float = Field(
+        30, ge=0, alias=c.SQRL_AUTH_TOKEN_EXPIRE_MINUTES, 
+        description="Expiration time for access tokens in minutes"
+    )
+    auth_credential_origins: list[str] = Field(
+        ["https://squirrels-analytics.github.io"], alias=c.SQRL_AUTH_CREDENTIAL_ORIGINS, 
+        description="Allowed origins for credentials (cookies)"
+    )
+
+    # Permissions
+    elevated_access_level: Literal["admin", "member", "guest"] = Field(
+        "admin", alias=c.SQRL_PERMISSIONS_ELEVATED_ACCESS_LEVEL, 
+        description="Minimum access level to access the studio"
+    )
+
+    # Parameters Cache
+    parameters_cache_size: int = Field(
+        1024, ge=0, alias=c.SQRL_PARAMETERS_CACHE_SIZE, 
+        description="Cache size for parameter configs"
+    )
+    parameters_cache_ttl_minutes: float = Field(
+        60, gt=0, alias=c.SQRL_PARAMETERS_CACHE_TTL_MINUTES, 
+        description="Cache TTL for parameter configs in minutes"
+    )
+    parameters_datasource_refresh_minutes: float = Field(
+        60, alias=c.SQRL_PARAMETERS_DATASOURCE_REFRESH_MINUTES, 
+        description="Interval in minutes for refreshing data sources. A non-positive value disables auto-refresh"
+    )
+    
+    # Datasets Cache
+    datasets_cache_size: int = Field(
+        128, ge=0, alias=c.SQRL_DATASETS_CACHE_SIZE, 
+        description="Cache size for dataset results"
+    )
+    datasets_cache_ttl_minutes: float = Field(
+        60, gt=0, alias=c.SQRL_DATASETS_CACHE_TTL_MINUTES, 
+        description="Cache TTL for dataset results in minutes"
+    )
+    datasets_max_rows_for_ai: int = Field(
+        100, ge=0, alias=c.SQRL_DATASETS_MAX_ROWS_FOR_AI, 
+        description="Max rows for AI queries"
+    )
+    datasets_max_rows_output: int = Field(
+        100000, ge=0, alias=c.SQRL_DATASETS_MAX_ROWS_OUTPUT, 
+        description="Max rows for dataset output"
+    )
+    datasets_sql_timeout_seconds: float = Field(
+        2.0, gt=0, alias=c.SQRL_DATASETS_SQL_TIMEOUT_SECONDS, 
+        description="Timeout for SQL operations in seconds"
+    )
+
+    # Dashboards Cache
+    dashboards_cache_size: int = Field(
+        128, ge=0, alias=c.SQRL_DASHBOARDS_CACHE_SIZE, 
+        description="Cache size for dashboards"
+    )
+    dashboards_cache_ttl_minutes: float = Field(
+        60, gt=0, alias=c.SQRL_DASHBOARDS_CACHE_TTL_MINUTES, 
+        description="Cache TTL for dashboards in minutes"
+    )
+    
+    # Seeds
+    seeds_infer_schema: bool = Field(
+        True, alias=c.SQRL_SEEDS_INFER_SCHEMA, 
+        description="Whether to infer schema for seeds"
+    )
+    seeds_na_values: list[str] = Field(
+        ["NA"], alias=c.SQRL_SEEDS_NA_VALUES, 
+        description="List of N/A values for seeds"
+    )
+
+    # Connections
+    connections_default_name_used: str = Field(
+        "default", alias=c.SQRL_CONNECTIONS_DEFAULT_NAME_USED, 
+        description="Default connection name to use"
+    )
+
+    # VDL
+    vdl_catalog_db_path: str = Field(
+        "ducklake:{project_path}/target/vdl_catalog.duckdb", alias=c.SQRL_VDL_CATALOG_DB_PATH, 
+        description="Path to the DuckDB catalog database"
+    )
+    vdl_data_path: str = Field(
+        "{project_path}/target/vdl_data/", alias=c.SQRL_VDL_DATA_PATH, 
+        description="Path to the VDL data directory"
+    )
+
+    # Studio
+    studio_base_url: str = Field(
+        "https://squirrels-analytics.github.io/squirrels-studio-v2", alias=c.SQRL_STUDIO_BASE_URL, 
+        description="Base URL for Squirrels Studio"
+    )
+
+    # Logging
+    logging_level: str = Field(
+        "INFO", alias=c.SQRL_LOGGING_LEVEL, 
+        description="Logging level"
+    )
+    logging_format: str = Field(
+        "text", alias=c.SQRL_LOGGING_FORMAT, 
+        description="Logging format"
+    )
+    logging_to_file: bool | str = Field(
+        False, alias=c.SQRL_LOGGING_TO_FILE, 
+        description="Whether to log to file. Can be set to true to use the default 'logs/' folder, or a folder path to write to a custom folder."
+    )
+    logging_file_size_mb: float = Field(
+        50, gt=0, alias=c.SQRL_LOGGING_FILE_SIZE_MB, 
+        description="Max log file size in MB"
+    )
+    logging_file_backup_count: int = Field(
+        1, ge=0, alias=c.SQRL_LOGGING_FILE_BACKUP_COUNT, 
+        description="Number of backup log files to keep"
+    )
+    
+    @field_validator("project_path")
+    @classmethod
+    def validate_project_path_exists(cls, v: str) -> str:
+        """Validate that the project_path is a folder that contains a squirrels.yml file."""
+        path = Path(v)
+        if not path.exists():
+            raise ValueError(f"Project path does not exist: {v}")
+        if not path.is_dir():
+            raise ValueError(f"Project path must be a directory, not a file: {v}")
+        # squirrels_yml = path / c.MANIFEST_FILE
+        # if not squirrels_yml.exists():
+        #     raise ValueError(f"Project path must contain a {c.MANIFEST_FILE} file: {v}")
+        return v
+    
+    @field_validator("auth_credential_origins", mode="before")
+    @classmethod
+    def parse_origins(cls, v: Any) -> list[str]:
+        if isinstance(v, str):
+            res = [x.strip() for x in v.split(",") if x.strip()]
+            return res or ["https://squirrels-analytics.github.io"]
+        return v
+
+    @field_validator("seeds_na_values", mode="before")
+    @classmethod
+    def parse_json_list(cls, v: Any) -> list[str]:
+        if isinstance(v, str):
+            try:
+                parsed = json.loads(v)
+                if not isinstance(parsed, list):
+                    raise ValueError(f"The {c.SQRL_SEEDS_NA_VALUES} environment variable must be a JSON list")
+                return parsed
+            except json.JSONDecodeError:
+                return []
+        return v
+    
+    @field_validator("logging_to_file", mode="before")
+    @classmethod
+    def parse_logging_to_file(cls, v: Any) -> bool | str:
+        if isinstance(v, str):
+            v_lower = v.lower()
+            if v_lower in ("true", "t", "1", "yes", "y", "on"):
+                return True
+            if v_lower in ("false", "f", "0", "no", "n", "off"):
+                return False
+            return v
+        return bool(v)
+
+    @field_validator("seeds_infer_schema", mode="before")
+    @classmethod
+    def parse_bool(cls, v: Any) -> bool:
+        if isinstance(v, str):
+            return v.lower() in ("true", "t", "1", "yes", "y", "on")
+        return bool(v)
+
+    @model_validator(mode="after")
+    def format_paths_with_filepath(self) -> Self:
+        """Format paths containing {filepath} placeholder with the actual filepath value."""
+        self.auth_db_file_path = self.auth_db_file_path.format(project_path=self.project_path)
+        self.vdl_catalog_db_path = self.vdl_catalog_db_path.format(project_path=self.project_path)
+        self.vdl_data_path = self.vdl_data_path.format(project_path=self.project_path)
+        return self
+    

squirrels/_exceptions.py

@@ -2,44 +2,16 @@ class InvalidInputError(Exception):
     """
     Use this exception when the error is due to providing invalid inputs to the REST API
 
-    Specific error code ranges are reserved for specific categories of errors.
-    0-19: 401 unauthorized errors
-    20-39: 403 forbidden errors
-    40-59: 404 not found errors
-    60-69: 409 conflict errors
-    70-99: Reserved for future use
-    100-199: 400 bad request errors related to authentication
-    200-299: 400 bad request errors related to data analytics
-
-    Error code definitions:
-    0 - Incorrect username or password
-    1 - Invalid authorization token
-    2 - Username not found for password change
-    3 - Incorrect password for password change
-    20 - Authorized user is forbidden to add or update users
-    21 - Authorized user is forbidden to delete users
-    22 - Cannot delete your own user
-    23 - Cannot delete the admin user
-    24 - Setting the admin user to non-admin is not permitted
-    25 - User does not have permission to access the dataset / dashboard
-    26 - User does not have permission to build the virtual data environment
-    27 - User does not have permission to query data models
-    40 - No token found for token_id
-    41 - No user found for username
-    60 - An existing build process is already running and a concurrent build is not allowed
-    61 - Model depends on static data models that cannot be found
-    100 - Missing required field 'username' or 'password' when adding a new user
-    101 - Username already exists when adding a new user
-    102 - Invalid user data when adding a new user
-    200 - Invalid value for dataset parameter
-    201 - Invalid query parameter provided
-    202 - Could not determine parent parameter for parameter refresh
-    203 - SQL query must be provided
-    204 - Failed to run provided SQL query
+    Attributes:
+        status_code: The HTTP status code to return
+        error: A short error message that should never change in the future
+        error_description: A detailed error message (that is allowed to change in the future)
     """
-    def __init__(self, error_code: int, message: str, *args) -> None:
-        self.error_code = error_code
-        super().__init__(message, *args)
+    def __init__(self, status_code: int, error: str, error_description: str, *args) -> None:
+        self.status_code = status_code
+        self.error = error
+        self.error_description = error_description
+        super().__init__(error_description, *args)
 
 
 class ConfigurationError(Exception):

squirrels/_http_error_responses.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+from starlette.requests import Request
+from starlette.responses import JSONResponse
+
+from ._exceptions import InvalidInputError
+
+
+def _strip_path_suffix_from_base_url(request: Request, *, strip_path_suffix: str | None) -> str:
+    """
+    Return a base URL with a known mount suffix removed.
+
+    In Squirrels, the main app mounts sub-apps like `/api/0` and `/mcp`. When we're
+    handling a request inside those sub-apps, `request.base_url` includes the mount
+    path. For `WWW-Authenticate` we want to point to a top-level endpoint on the main
+    app (same mount as the main app itself), so we strip only the *sub-app mount*
+    suffix (e.g. `/api/0` or `/mcp`), not any outer mount path.
+    """
+    base_url = str(request.base_url).rstrip("/")
+    if strip_path_suffix:
+        suffix = strip_path_suffix.rstrip("/")
+        if suffix and base_url.endswith(suffix):
+            base_url = base_url[: -len(suffix)]
+    return base_url.rstrip("/")
+
+
+def invalid_input_error_to_json_response(
+    request: Request,
+    exc: InvalidInputError,
+    *,
+    oauth_resource_metadata_path: str = "/.well-known/oauth-protected-resource",
+    strip_path_suffix: str | None = None,
+    is_mcp: bool = False,
+) -> JSONResponse:
+    """
+    Convert an InvalidInputError into the standard Squirrels JSON error response.
+
+    For 401s, also sets `WWW-Authenticate` with a top-level `resource_metadata` URL.
+    """
+    response = JSONResponse(
+        status_code=exc.status_code,
+        content={"error": exc.error, "error_description": exc.error_description},
+    )
+
+    if exc.status_code == 401:
+        top_level_base_url = _strip_path_suffix_from_base_url(request, strip_path_suffix=strip_path_suffix)
+        resource_metadata_url = f"{top_level_base_url}{oauth_resource_metadata_path}"
+        realm = "mcp" if is_mcp else "api"
+        response.headers["WWW-Authenticate"] = f'Bearer realm="{realm}", resource_metadata="{resource_metadata_url}"'
+
+    return response
+

squirrels/_initializer.py

@@ -67,8 +67,8 @@ class Initializer:
     def _copy_federate_file(self, filepath: str):
         self._copy_file(Path(c.MODELS_FOLDER, c.FEDERATES_FOLDER, filepath))
 
-    def _copy_database_file(self, filepath: str):
-        self._copy_file(Path(c.DATABASE_FOLDER, filepath))
+    def _copy_resource_file(self, filepath: str):
+        self._copy_file(Path(c.RESOURCES_FOLDER, filepath))
     
     def _copy_pyconfig_file(self, filepath: str):
         self._copy_file(Path(c.PYCONFIGS_FOLDER, filepath))
@@ -123,7 +123,7 @@ class Initializer:
         CONNECTIONS, PARAMETERS, BUILD, FEDERATE, DASHBOARD, ADMIN_PASSWORD = options
 
         # Add project name prompt if not provided
-        if self.project_name is None and not args.curr_dir:
+        if self.project_name is None and not self.use_curr_dir:
             questions = [
                 inquirer.Text('project_name', message="What is your project folder name? (leave blank to create in current directory)")
             ]
@@ -264,9 +264,10 @@ class Initializer:
             self._copy_dashboard_file(c.DASHBOARD_FILE_STEM + ".py")
             self._copy_dashboard_file(c.DASHBOARD_FILE_STEM + ".yml")
         
-        self._copy_database_file(c.EXPENSES_DB)
+        self._copy_resource_file(c.EXPENSES_DB)
+        self._copy_file(Path(c.RESOURCES_FOLDER, c.PUBLIC_FOLDER, ".gitkeep"))
         
-        print(f"\nSuccessfully created new Squirrels project in current directory!\n")
+        print(f"\nSuccessfully created new Squirrels project!\n")
     
     def get_file(self, args):
         if args.file_name == c.DOTENV_FILE:
@@ -309,7 +310,7 @@ class Initializer:
             self._copy_dashboard_file(args.file_name + ".py")
             self._copy_dashboard_file(args.file_name + ".yml")
         elif args.file_name in (c.EXPENSES_DB, c.WEATHER_DB):
-            self._copy_database_file(args.file_name)
+            self._copy_resource_file(args.file_name)
         elif args.file_name in (c.SEED_CATEGORY_FILE_STEM, c.SEED_SUBCATEGORY_FILE_STEM):
             self._copy_seed_file(args.file_name + ".csv")
             self._copy_seed_file(args.file_name + ".yml")

squirrels/_logging.py

@@ -0,0 +1,121 @@
+from pathlib import Path
+from logging.handlers import RotatingFileHandler
+from uuid import uuid4
+import logging as l, json
+
+from . import _constants as c, _utils as u
+from ._request_context import get_request_id
+
+
+class _BaseFormatter(l.Formatter):
+    def _format_helper(self, level_for_print: str, record: l.LogRecord) -> str:
+        # Save original levelname
+        original_levelname = record.levelname
+
+        # Add padding to the levelname for printing
+        visible_length = len(record.levelname) + 1
+        padding_needed = max(1, 9 - visible_length)
+        padded_level = f"{level_for_print}:{' ' * padding_needed}"
+        record.levelname = padded_level
+        
+        # Format the message
+        formatted = super().format(record)
+        
+        # Append request ID if available
+        request_id = get_request_id()
+        request_id_str = f" [req_id: {request_id}]" if request_id else ""
+        formatted = formatted.replace("{request_id}", request_id_str)
+        
+        # Restore original levelname
+        record.levelname = original_levelname
+        
+        return formatted
+
+
+class _ColoredFormatter(_BaseFormatter):
+    """Custom formatter that adds colors to log levels for terminal output"""
+    
+    # ANSI color codes
+    COLORS = {
+        'DEBUG': '\033[36m',      # Cyan
+        'INFO': '\033[32m',       # Green
+        'WARNING': '\033[33m',    # Yellow
+        'ERROR': '\033[31m',      # Red
+        'CRITICAL': '\033[35m',   # Magenta
+    }
+    RESET = '\033[0m'
+    BOLD = '\033[1m'
+    
+    def format(self, record: l.LogRecord) -> str:
+        # Add color to levelname with colon and padding
+        color = self.COLORS.get(record.levelname, '')
+        colored_level = f"{color}{record.levelname}{self.RESET}"
+        return self._format_helper(colored_level, record)
+
+
+class _PlainFormatter(_BaseFormatter):
+    """Custom formatter that adds colon to log levels for file output"""
+    
+    def format(self, record: l.LogRecord) -> str:
+        return self._format_helper(record.levelname, record)
+
+
+class _CustomJsonFormatter(l.Formatter):
+    def format(self, record: l.LogRecord) -> str:
+        super().format(record)
+        request_id = get_request_id()
+        info = {
+            "timestamp": self.formatTime(record),
+            "level": record.levelname,
+            "message": record.getMessage(),
+            "request_id": request_id,
+        }
+        output = {
+            "data": record.__dict__.get("data", {}),
+            "info": info
+        }
+        return json.dumps(output)
+
+
+def get_logger(
+    project_path: str, log_to_file: bool | str, log_level: str, log_format: str, log_file_size_mb: float, log_file_backup_count: int
+) -> u.Logger:
+    logger = u.Logger(name=uuid4().hex, level=log_level.upper())
+
+    # Determine the formatter based on log_format
+    if log_format.lower() == "json":
+        stdout_formatter = _CustomJsonFormatter()
+        file_formatter = _CustomJsonFormatter()
+    elif log_format.lower() == "text":
+        # Use colored formatter for stdout, plain formatter with colon for file
+        format_string = "%(levelname)s [%(asctime)s]{request_id} %(message)s"
+        stdout_formatter = _ColoredFormatter(format_string, datefmt="%Y-%m-%d %H:%M:%S")
+        file_formatter = _PlainFormatter(format_string, datefmt="%Y-%m-%d %H:%M:%S")
+    else:
+        raise ValueError("log_format must be either 'text' or 'json'")
+    
+    if log_to_file:
+        if isinstance(log_to_file, str):
+            log_file_path = Path(project_path, log_to_file, c.LOGS_FILE)
+        else:
+            log_file_path = Path(project_path, c.LOGS_FOLDER, c.LOGS_FILE)
+        
+        log_file_path.parent.mkdir(parents=True, exist_ok=True)
+
+        # Rotating file handler
+        file_handler = RotatingFileHandler(
+            log_file_path,
+            maxBytes=int(log_file_size_mb * 1024 * 1024),
+            backupCount=log_file_backup_count
+        )
+        file_handler.setLevel(log_level.upper())
+        file_handler.setFormatter(file_formatter)
+        logger.addHandler(file_handler)
+    
+    else:
+        stdout_handler = l.StreamHandler()
+        stdout_handler.setLevel(log_level.upper())
+        stdout_handler.setFormatter(stdout_formatter)
+        logger.addHandler(stdout_handler)
+    
+    return logger