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

squirrels/_connection_set.py

@@ -1,11 +1,11 @@
-from dataclasses import dataclass
-from sqlalchemy import Engine, create_engine
-import time, pandas as pd
+from typing import Any
+from dataclasses import dataclass, field
+from sqlalchemy import Engine
+import time, polars as pl
 
 from . import _utils as u, _constants as c, _py_module as pm
 from .arguments.init_time_args import ConnectionsArgs
-from ._environcfg import EnvironConfig
-from ._manifest import ManifestConfig
+from ._manifest import ManifestConfig, ConnectionProperties, ConnectionType
 
 
 @dataclass
@@ -16,44 +16,56 @@ class ConnectionSet:
     Attributes:
         _engines: A dictionary of connection name to the corresponding sqlalchemy engine
     """
-    _engines: dict[str, Engine]
+    _connections: dict[str, ConnectionProperties | Any] = field(default_factory=dict)
 
-    def get_engines_as_dict(self):
-        return self._engines.copy()
+    def get_connections_as_dict(self):
+        return self._connections.copy()
     
-    def _get_engine(self, conn_name: str) -> Engine:
+    def get_connection(self, conn_name: str) -> ConnectionProperties | Any:
         try:
-            connection_pool = self._engines[conn_name]
+            connection = self._connections[conn_name]
         except KeyError as e:
             raise u.ConfigurationError(f'Connection name "{conn_name}" was not configured') from e
-        return connection_pool
+        return connection
     
-    def run_sql_query_from_conn_name(self, query: str, conn_name: str, placeholders: dict = {}) -> pd.DataFrame:
-        engine = self._get_engine(conn_name)
+    def run_sql_query_from_conn_name(self, query: str, conn_name: str, placeholders: dict = {}) -> pl.DataFrame:
+        conn = self.get_connection(conn_name)
         try:
-            df = pd.read_sql(query, engine, params=placeholders)
+            if isinstance(conn, ConnectionProperties) and (conn.type == ConnectionType.CONNECTORX or conn.type == ConnectionType.ADBC):
+                if len(placeholders) > 0:
+                    raise u.ConfigurationError(f"Connection '{conn_name}' is a ConnectorX or ADBC connection, which does not support placeholders")
+                df = pl.read_database_uri(query, conn.uri, engine=conn.type.value)
+            elif isinstance(conn, ConnectionProperties) and conn.type == ConnectionType.SQLALCHEMY:
+                with conn.engine.connect() as connection:
+                    df = pl.read_database(query, connection, execute_options={"parameters": placeholders})
+            else:
+                df = pl.read_database(query, conn, execute_options={"parameters": placeholders}) # type: ignore
             return df
         except Exception as e:
             raise RuntimeError(e) from e
 
     def dispose(self) -> None:
         """
-        Disposes of all the engines in this ConnectionSet
+        Disposes / closes all the connections in this ConnectionSet
         """
-        for pool in self._engines.values():
-            if isinstance(pool, Engine):
-                pool.dispose()
+        for conn in self._connections.values():
+            if isinstance(conn, Engine):
+                conn.dispose()
+            elif isinstance(conn, ConnectionProperties):
+                if conn.type == ConnectionType.SQLALCHEMY:
+                    conn.engine.dispose()
+            elif hasattr(conn, 'close'):
+                conn.close()
 
 
 class ConnectionSetIO:
 
     @classmethod
-    def load_conn_py_args(cls, logger: u.Logger, env_cfg: EnvironConfig, manifest_cfg: ManifestConfig) -> ConnectionsArgs:
+    def load_conn_py_args(cls, logger: u.Logger, base_path: str, env_vars: dict[str, str], manifest_cfg: ManifestConfig) -> ConnectionsArgs:
         start = time.time()
         
         proj_vars = manifest_cfg.project_variables.model_dump()
-        env_vars = env_cfg.get_all_env_vars()
-        conn_args = ConnectionsArgs(proj_vars, env_vars, env_cfg.get_credential)
+        conn_args = ConnectionsArgs(base_path, proj_vars, env_vars)
         
         logger.log_activity_time("setting up arguments for connections.py", start)
         return conn_args
@@ -68,13 +80,14 @@ class ConnectionSetIO:
             A ConnectionSet with the DB connections from both squirrels.yml and connections.py
         """
         start = time.time()
-        engines: dict[str, Engine] = {}
+        connections: dict[str, ConnectionProperties | Any] = {}
         
         for config in manifest_cfg.connections.values():
-            engines[config.name] = create_engine(config.url)
+            connections[config.name] = ConnectionProperties(label=config.label, type=config.type, uri=config.uri)
 
-        pm.run_pyconfig_main(base_path, c.CONNECTIONS_FILE, {"connections": engines, "sqrl": conn_args})
-        conn_set = ConnectionSet(engines)
+        pm.run_pyconfig_main(base_path, c.CONNECTIONS_FILE, {"connections": connections, "sqrl": conn_args})
+
+        conn_set = ConnectionSet(connections)
 
         logger.log_activity_time("creating sqlalchemy engines", start)
         return conn_set

squirrels/_constants.py

@@ -1,25 +1,51 @@
 # Squirrels CLI commands
-INIT_CMD = 'init'
+INIT_CMD = 'new'
 GET_FILE_CMD = 'get-file'
 DEPS_CMD = 'deps'
+BUILD_CMD = 'build'
 COMPILE_CMD = 'compile'
 RUN_CMD = 'run'
+DUCKDB_CMD = 'duckdb'
 
-# Environment config keys
-JWT_SECRET_KEY = 'jwt_secret'
+# Environment variables
+SQRL_SECRET_KEY = 'SQRL_SECRET__KEY'
+SQRL_SECRET_ADMIN_PASSWORD = 'SQRL_SECRET__ADMIN_PASSWORD'
+
+SQRL_AUTH_DB_FILE_PATH = 'SQRL_AUTH__DB_FILE_PATH'
+SQRL_AUTH_TOKEN_EXPIRE_MINUTES = 'SQRL_AUTH__TOKEN_EXPIRE_MINUTES'
+
+SQRL_PARAMETERS_CACHE_SIZE = 'SQRL_PARAMETERS__CACHE_SIZE'
+SQRL_PARAMETERS_CACHE_TTL_MINUTES = 'SQRL_PARAMETERS__CACHE_TTL_MINUTES'
+
+SQRL_DATASETS_CACHE_SIZE = 'SQRL_DATASETS__CACHE_SIZE'
+SQRL_DATASETS_CACHE_TTL_MINUTES = 'SQRL_DATASETS__CACHE_TTL_MINUTES'
+
+SQRL_DASHBOARDS_CACHE_SIZE = 'SQRL_DASHBOARDS__CACHE_SIZE'
+SQRL_DASHBOARDS_CACHE_TTL_MINUTES = 'SQRL_DASHBOARDS__CACHE_TTL_MINUTES'
+
+SQRL_SEEDS_INFER_SCHEMA = 'SQRL_SEEDS__INFER_SCHEMA'
+SQRL_SEEDS_NA_VALUES = 'SQRL_SEEDS__NA_VALUES'
+
+SQRL_TEST_SETS_DEFAULT_NAME_USED = 'SQRL_TEST_SETS__DEFAULT_NAME_USED'
+
+SQRL_CONNECTIONS_DEFAULT_NAME_USED = 'SQRL_CONNECTIONS__DEFAULT_NAME_USED'
+
+SQRL_DUCKDB_VENV_DB_FILE_PATH = 'SQRL_DUCKDB_VENV__DB_FILE_PATH'
 
 # Folder/File names
 PACKAGE_DATA_FOLDER = 'package_data'
 BASE_PROJECT_FOLDER = 'base_project'
-ASSETS_FOLDER = 'assets'
-TEMPLATES_FOLDER = 'templates'
 
-ENV_CONFIG_FILE = 'env.yml'
+GLOBAL_ENV_FOLDER = '.squirrels'
 MANIFEST_JINJA_FILE = 'squirrels.yml.j2'
 CONNECTIONS_YML_FILE = 'connections.yml'
 PARAMETERS_YML_FILE = 'parameters.yml'
 DASHBOARDS_YML_FILE = 'dashboards.yml'
 MANIFEST_FILE = 'squirrels.yml'
+DUCKDB_INIT_FILE = 'duckdb_init.sql'
+DOTENV_FILE = '.env'
+DOTENV_LOCAL_FILE = '.env.local'
+GITIGNORE_FILE = '.gitignore'
 
 LOGS_FOLDER = 'logs'
 LOGS_FILE = 'squirrels.log'
@@ -28,8 +54,12 @@ DATABASE_FOLDER = 'assets'
 PACKAGES_FOLDER = 'sqrl_packages'
 
 MACROS_FOLDER = 'macros'
+MACROS_FILE = 'macros_example.sql'
 
 MODELS_FOLDER = 'models'
+SOURCES_FILE = 'sources.yml'
+BUILDS_FOLDER = 'builds'
+BUILD_FILE_STEM = 'build_example'
 DBVIEWS_FOLDER = 'dbviews'
 DBVIEW_FILE_STEM = 'dbview_example'
 FEDERATES_FOLDER = 'federates'
@@ -43,35 +73,17 @@ AUTH_FILE = 'auth.py'
 CONNECTIONS_FILE = 'connections.py'
 CONTEXT_FILE = 'context.py'
 PARAMETERS_FILE = 'parameters.py'
+USER_FILE = 'user.py'
+ADMIN_USERNAME = 'admin'
 
 TARGET_FOLDER = 'target'
+DB_FILE = 'auth.sqlite'
 COMPILE_FOLDER = 'compile'
+DUCKDB_VENV_FILE = 'venv.duckdb'
 
 SEEDS_FOLDER = 'seeds'
-CATEGORY_SEED_FILE = 'seed_categories.csv'
-SUBCATEGORY_SEED_FILE = 'seed_subcategories.csv'
-
-# Dataset setting names
-AUTH_TOKEN_EXPIRE_SETTING = 'auth.token.expire_minutes'
-PARAMETERS_CACHE_SIZE_SETTING = 'parameters.cache.size'
-PARAMETERS_CACHE_TTL_SETTING = 'parameters.cache.ttl_minutes'
-RESULTS_CACHE_SIZE_SETTING = 'results.cache.size' # deprecated
-RESULTS_CACHE_TTL_SETTING = 'results.cache.ttl_minutes' # deprecated
-DATASETS_CACHE_SIZE_SETTING = 'datasets.cache.size'
-DATASETS_CACHE_TTL_SETTING = 'datasets.cache.ttl_minutes'
-DASHBOARDS_CACHE_SIZE_SETTING = 'dashboards.cache.size'
-DASHBOARDS_CACHE_TTL_SETTING = 'dashboards.cache.ttl_minutes'
-TEST_SET_DEFAULT_USED_SETTING = 'selection_test_sets.default_name_used'
-DEFAULT_TEST_SET_NAME = 'default'
-DB_CONN_DEFAULT_USED_SETTING = 'connections.default_name_used'
-DEFAULT_DB_CONN = 'default'
-DEFAULT_MATERIALIZE_SETTING = 'defaults.federates.materialized'
-DEFAULT_MATERIALIZE = 'table'
-SEEDS_INFER_SCHEMA_SETTING = 'seeds.infer_schema'
-SEEDS_NA_VALUES_SETTING = 'seeds.na_values'
-IN_MEMORY_DB_SETTING = 'in_memory_database'
-SQLITE = 'sqlite'
-DUCKDB = 'duckdb'
+SEED_CATEGORY_FILE_STEM = 'seed_categories'
+SEED_SUBCATEGORY_FILE_STEM = 'seed_subcategories'
 
 # Selection cfg sections
 USER_ATTRIBUTES_SECTION = 'user_attributes'
@@ -97,10 +109,8 @@ PNG = "png"
 HTML = "html"
 
 # Function names
-GET_USER_FUNC = "get_user_if_valid"
-DEP_FUNC = "dependencies"
 MAIN_FUNC = "main"
 
 # Regex
-date_regex = r"^\d{4}\-\d{2}\-\d{2}$"
-color_regex = r"^#[0-9a-fA-F]{6}$"
+DATE_REGEX = r"^\d{4}\-\d{2}\-\d{2}$"
+COLOR_REGEX = r"^#[0-9a-fA-F]{6}$"

squirrels/_dashboards_io.py

@@ -1,18 +1,37 @@
 from typing import Type, TypeVar, Callable, Coroutine, Any
+from enum import Enum
 from dataclasses import dataclass
-import inspect, os, time
+from pydantic import BaseModel, Field
+import os, time
 
 from .arguments.run_time_args import DashboardArgs
 from ._py_module import PyModule
+from ._manifest import AnalyticsOutputConfig
+from ._exceptions import InvalidInputError, ConfigurationError, FileExecutionError
 from . import _constants as c, _utils as u, dashboards as d
 
 T = TypeVar('T', bound=d.Dashboard)
 
 
+class DashboardFormat(Enum):
+    PNG = "png"
+    HTML = "html"
+
+class DashboardDependencies(BaseModel):
+    name: str
+    dataset: str
+    fixed_parameters: list[dict[str, str]] = Field(default_factory=list)
+
+class DashboardConfig(AnalyticsOutputConfig):
+    format: DashboardFormat = Field(default=DashboardFormat.PNG)
+    depends_on: list[DashboardDependencies] = Field(default_factory=list)
+
+
 @dataclass
-class DashboardFunction:
+class DashboardDefinition:
     dashboard_name: str
     filepath: str
+    config: DashboardConfig
 
     @property
     def dashboard_func(self) -> Callable[[DashboardArgs], Coroutine[Any, Any, d.Dashboard]]:
@@ -22,23 +41,16 @@ class DashboardFunction:
         return self._dashboard_func
 
     def get_dashboard_format(self) -> str:
-        return_type = inspect.signature(self.dashboard_func).return_annotation
-        assert issubclass(return_type, d.Dashboard), f"Function must return Dashboard type"
-        if return_type == d.PngDashboard:
-            return c.PNG
-        elif return_type == d.HtmlDashboard:
-            return c.HTML
-        else:
-            raise NotImplementedError(f"Dashboard format {return_type} not supported")
+        return self.config.format.value
     
     async def get_dashboard(self, args: DashboardArgs, *, dashboard_type: Type[T] = d.Dashboard) -> T:
         try:
             dashboard = await self.dashboard_func(args)
             assert isinstance(dashboard, dashboard_type), f"Function does not return expected Dashboard type: {dashboard_type}"
-        except (u.InvalidInputError, u.ConfigurationError, u.FileExecutionError) as e:
+        except (InvalidInputError, ConfigurationError, FileExecutionError) as e:
             raise e
         except Exception as e:
-            raise u.FileExecutionError(f'Failed to run "{c.MAIN_FUNC}" function for dashboard "{self.dashboard_name}"', e) from e
+            raise FileExecutionError(f'Failed to run "{c.MAIN_FUNC}" function for dashboard "{self.dashboard_name}"', e) from e
 
         return dashboard
 
@@ -46,7 +58,7 @@ class DashboardFunction:
 class DashboardsIO:
 
     @classmethod
-    def load_files(cls, logger: u.Logger, base_path: str) -> dict[str, DashboardFunction]:
+    def load_files(cls, logger: u.Logger, base_path: str) -> dict[str, DashboardDefinition]:
         start = time.time()
         
         dashboards_by_name = {}
@@ -54,8 +66,14 @@ class DashboardsIO:
             for file in filenames:
                 filepath = os.path.join(dp, file)
                 file_stem, extension = os.path.splitext(file)
-                if extension == '.py':
-                    dashboards_by_name[file_stem] = DashboardFunction(file_stem, filepath)
-        
+                if not extension == '.py':
+                    continue
+                
+                # Check for corresponding .yml file
+                yml_path = os.path.join(dp, file_stem + '.yml')
+                config_dict = u.load_yaml_config(yml_path) if os.path.exists(yml_path) else {}
+                config = DashboardConfig(name=file_stem, **config_dict)
+                dashboards_by_name[file_stem] = DashboardDefinition(file_stem, filepath, config)
+                
         logger.log_activity_time("loading files for dashboards", start)
         return dashboards_by_name

squirrels/_exceptions.py

@@ -0,0 +1,57 @@
+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 - Cannot change the admin user
+    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
+    """
+    def __init__(self, error_code: int, message: str, *args) -> None:
+        self.error_code = error_code
+        super().__init__(message, *args)
+
+
+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