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

squirrels/_manifest.py

@@ -1,22 +1,32 @@
-from typing import Any
+from functools import cached_property
+from typing import Literal, Any
+from urllib.parse import urlparse
+from sqlalchemy import Engine, create_engine
 from typing_extensions import Self
 from enum import Enum
 from pydantic import BaseModel, Field, field_validator, model_validator, ValidationInfo, ValidationError
-import yaml, time
+import yaml, time, re
 
-from . import _constants as c, _utils as _u
-from ._environcfg import EnvironConfig
+from . import _constants as c, _utils as u
 
 
 class ProjectVarsConfig(BaseModel, extra="allow"):
     name: str
     label: str = ""
+    description: str = ""
     major_version: int
 
+    @field_validator("name")
+    @classmethod
+    def validate_name(cls, v: str) -> str:
+        if not re.fullmatch(r"[A-Za-z0-9_-]+", v):
+            raise ValueError("Project name must only contain alphanumeric characters, underscores, and dashes.")
+        return v
+
     @model_validator(mode="after")
     def finalize_label(self) -> Self:
         if self.label == "":
-            self.label = self.name
+            self.label = u.to_title_case(self.name)
         return self
 
 
@@ -36,41 +46,133 @@ class _ConfigWithNameBaseModel(BaseModel):
     name: str
 
 
-class DbConnConfig(_ConfigWithNameBaseModel):
-    credential: str | None = None
-    url: str
+class ConnectionTypeEnum(Enum):
+    SQLALCHEMY = "sqlalchemy"
+    CONNECTORX = "connectorx"
+    ADBC = "adbc"
+    DUCKDB = "duckdb"
 
-    def finalize_url(self, base_path: str, env_cfg: EnvironConfig) -> Self:
-        username, password = env_cfg.get_credential(self.credential)
-        self.url = self.url.format(username=username, password=password, project_path=base_path)
+
+class ConnectionProperties(BaseModel):
+    """
+    A class for holding the properties of a connection
+
+    Arguments:
+        type: The type of connection, one of "sqlalchemy", "connectorx", or "adbc"
+        uri: The URI for the connection
+    """
+    label: str | None = None
+    type: ConnectionTypeEnum = Field(default=ConnectionTypeEnum.SQLALCHEMY)
+    uri: str
+    sa_create_engine_args: dict[str, Any] = Field(default_factory=dict)
+
+    @cached_property
+    def engine(self) -> Engine:
+        """
+        Creates and caches a SQLAlchemy engine if the connection type is sqlalchemy.
+        Returns None for other connection types.
+        """
+        if self.type == ConnectionTypeEnum.SQLALCHEMY:
+            return create_engine(self.uri, **self.sa_create_engine_args)
+        else:
+            raise ValueError(f'Connection type "{self.type}" does not support engine property')
+
+    @cached_property
+    def dialect(self) -> str:
+        default_dialect = None
+        if self.type == ConnectionTypeEnum.SQLALCHEMY:
+            dialect = self.engine.dialect.name
+        elif self.type == ConnectionTypeEnum.DUCKDB:
+            dialect = self.uri.split(':')[0]
+            default_dialect = 'duckdb'
+        else:
+            url = urlparse(self.uri)
+            dialect = url.scheme
+        
+        processed_dialect = next((d for d in ['sqlite', 'postgres', 'mysql', 'duckdb'] if dialect.lower().startswith(d)), default_dialect)
+        dialect = processed_dialect if processed_dialect is not None else dialect
+        return dialect
+    
+    @cached_property
+    def attach_uri_for_duckdb(self) -> str | None:
+        if self.type == ConnectionTypeEnum.DUCKDB:
+            return self.uri
+        elif self.type == ConnectionTypeEnum.SQLALCHEMY:
+            url = self.engine.url
+            host = url.host
+            port = url.port
+            username = url.username
+            password = url.password
+            database = url.database
+            database_as_file = database if database is not None else ""
+        else:
+            url = urlparse(self.uri)
+            host = url.hostname
+            port = url.port
+            username = url.username
+            password = url.password
+            database = url.path.lstrip('/')
+            database_as_file = self.uri.replace(f"{self.dialect}://", "")
+        
+        if self.dialect in ('postgres', 'mysql'):
+            attach_uri = f"{self.dialect}:dbname={database} user={username} password={password} host={host} port={port}"
+        elif self.dialect == "sqlite":
+            attach_uri = f"{self.dialect}:{database_as_file}"
+        elif self.dialect == "duckdb":
+            attach_uri = database_as_file
+        else:
+            attach_uri = None
+        
+        return attach_uri
+
+
+class DbConnConfig(ConnectionProperties, _ConfigWithNameBaseModel):
+    def finalize_uri(self, base_path: str) -> Self:
+        self.uri = self.uri.format(project_path=base_path)
         return self
 
 
+class DatasetConfigurablesConfig(BaseModel):
+    name: str
+    default: str
+
+
+class ConfigurablesConfig(DatasetConfigurablesConfig):
+    label: str = ""
+    description: str = ""
+
+
 class ParametersConfig(BaseModel):
     type: str
     factory: str
     arguments: dict[str, Any]
 
 
-class DbviewConfig(_ConfigWithNameBaseModel):
-    connection_name: str | None = None
+class PermissionScope(Enum):
+    PUBLIC = 0
+    PROTECTED = 1
+    PRIVATE = 2
 
 
-class FederateConfig(_ConfigWithNameBaseModel):
-    materialized: str | None = None
+class AuthenticationEnforcement(Enum):
+    REQUIRED = "required"
+    OPTIONAL = "optional"
+    DISABLED = "disabled"
 
+class AuthenticationType(Enum):
+    MANAGED = "managed"
+    EXTERNAL = "external"
 
-class DatasetScope(Enum):
-    PUBLIC = 0
-    PROTECTED = 1
-    PRIVATE = 2
+class AuthenticationConfig(BaseModel):
+    enforcement: AuthenticationEnforcement = AuthenticationEnforcement.OPTIONAL
+    type: AuthenticationType = AuthenticationType.MANAGED
 
 
 class AnalyticsOutputConfig(_ConfigWithNameBaseModel):
     label: str = ""
     description: str = ""
-    scope: DatasetScope = DatasetScope.PUBLIC
-    parameters: list[str] = Field(default_factory=list)
+    scope: PermissionScope = PermissionScope.PUBLIC
+    parameters: list[str] | None = Field(default=None, description="The list of parameter names used by the dataset/dashboard")
 
     @model_validator(mode="after")
     def finalize_label(self) -> Self:
@@ -80,19 +182,18 @@ class AnalyticsOutputConfig(_ConfigWithNameBaseModel):
 
     @field_validator("scope", mode="before")
     @classmethod
-    def validate_scope(cls, value: str, info: ValidationInfo) -> DatasetScope:
+    def validate_scope(cls, value: str, info: ValidationInfo) -> PermissionScope:
         try:
-            return DatasetScope[str(value).upper()]
+            return PermissionScope[str(value).upper()]
         except KeyError as e:
             name = info.data.get("name")
-            scope_list = [scope.name.lower() for scope in DatasetScope]
+            scope_list = [scope.name.lower() for scope in PermissionScope]
             raise ValueError(f'Scope "{value}" is invalid for dataset/dashboard "{name}". Scope must be one of {scope_list}') from e
 
 
 class DatasetConfig(AnalyticsOutputConfig):
     model: str = ""
-    traits: dict = Field(default_factory=dict)
-    default_test_set: str = ""
+    configurables: list[DatasetConfigurablesConfig] = Field(default_factory=list)
     
     def __hash__(self) -> int:
         return hash("dataset_"+self.name)
@@ -104,47 +205,27 @@ class DatasetConfig(AnalyticsOutputConfig):
         return self
 
 
-class DashboardConfig(AnalyticsOutputConfig):
-    def __hash__(self) -> int:
-        return hash("dashboard_"+self.name)
-
+class TestSetsUserConfig(BaseModel):
+    access_level: Literal["admin", "member", "guest"] = "guest"
+    custom_fields: dict[str, Any] = Field(default_factory=dict)
 
 class TestSetsConfig(_ConfigWithNameBaseModel):
-    datasets: list[str] | None = None
-    is_authenticated: bool = False
-    user_attributes: dict[str, Any] = Field(default_factory=dict)
+    user: TestSetsUserConfig = Field(default_factory=TestSetsUserConfig)
     parameters: dict[str, Any] = Field(default_factory=dict)
-
-    @model_validator(mode="after")
-    def finalize_is_authenticated(self) -> Self:
-        if len(self.user_attributes) > 0:
-            self.is_authenticated = True
-        return self
-
-
-class Settings(BaseModel):
-    data: dict[str, Any]
-    
-    def get_default_connection_name(self) -> str:
-        return self.data.get(c.DB_CONN_DEFAULT_USED_SETTING, c.DEFAULT_DB_CONN)
-    
-    def do_use_duckdb(self) -> bool:
-        return self.data.get(c.IN_MEMORY_DB_SETTING, c.SQLITE) == c.DUCKDB
+    configurables: dict[str, Any] = Field(default_factory=dict)
 
 
 class ManifestConfig(BaseModel):
-    env_cfg: EnvironConfig
     project_variables: ProjectVarsConfig
+    authentication: AuthenticationConfig = Field(default_factory=AuthenticationConfig)
     packages: list[PackageConfig] = Field(default_factory=list)
     connections: dict[str, DbConnConfig] = Field(default_factory=dict)
     parameters: list[ParametersConfig] = Field(default_factory=list)
+    configurables: dict[str, ConfigurablesConfig] = Field(default_factory=dict)
     selection_test_sets: dict[str, TestSetsConfig] = Field(default_factory=dict)
-    dbviews: dict[str, DbviewConfig] = Field(default_factory=dict)
-    federates: dict[str, FederateConfig] = Field(default_factory=dict)
     datasets: dict[str, DatasetConfig] = Field(default_factory=dict)
-    dashboards: dict[str, DashboardConfig] = Field(default_factory=dict)
-    settings: dict[str, Any] = Field(default_factory=dict)
     base_path: str = "."
+    env_vars: dict[str, str] = Field(default_factory=dict)
 
     @field_validator("packages")
     @classmethod
@@ -156,13 +237,13 @@ class ManifestConfig(BaseModel):
             set_of_directories.add(package.directory)
         return packages
     
-    @field_validator("connections", "selection_test_sets", "dbviews", "federates", "datasets", "dashboards", mode="before")
+    @field_validator("connections", "selection_test_sets", "datasets", "configurables", mode="before")
     @classmethod
     def names_are_unique(cls, values: list[dict] | dict[str, dict], info: ValidationInfo) -> dict[str, dict]:
         if isinstance(values, list):
             values_as_dict = {}
             for obj in values:
-                name = obj["name"]
+                name = u.normalize_name(obj["name"])
                 if name in values_as_dict:
                     raise ValueError(f'In the {info.field_name} section, the name "{name}" was specified multiple times')
                 values_as_dict[name] = obj
@@ -173,45 +254,93 @@ class ManifestConfig(BaseModel):
     @model_validator(mode="after")
     def finalize_connections(self) -> Self:
         for conn in self.connections.values():
-            conn.finalize_url(self.base_path, self.env_cfg)
+            conn.finalize_uri(self.base_path)
         return self
     
-    @property
-    def settings_obj(self) -> Settings:
-        return Settings(data=self.settings)
-
-    def get_default_test_set(self, dataset_name: str) -> TestSetsConfig:
+    @model_validator(mode="after")
+    def validate_authentication_and_scopes(self) -> Self:
+        """
+        Enforce authentication rules:
+        - If authentication.is_required, no dataset may be PUBLIC.
+        """
+        if self.authentication.enforcement == AuthenticationEnforcement.REQUIRED:
+            invalid = [name for name, ds in self.datasets.items() if ds.scope == PermissionScope.PUBLIC]
+            if invalid:
+                raise ValueError(
+                    "Authentication is required, so datasets cannot be public. "
+                    f"Update the scope for datasets: {invalid}"
+                )
+        return self
+    
+    @model_validator(mode="after")
+    def validate_dataset_configurables(self) -> Self:
+        """
+        Validate that dataset configurables reference valid project-level configurables.
+        """
+        for dataset_name, dataset_cfg in self.datasets.items():
+            for cfg_override in dataset_cfg.configurables:
+                if cfg_override.name not in self.configurables:
+                    raise ValueError(
+                        f'Dataset "{dataset_name}" references configurable "{cfg_override.name}" which is not defined '
+                        f'in the project configurables'
+                    )
+        return self
+    
+    def get_default_test_set(self) -> TestSetsConfig:
         """
         Raises KeyError if dataset name doesn't exist
         """
-        default_name_1 = self.datasets[dataset_name].default_test_set
-        default_name_2 = self.settings.get(c.TEST_SET_DEFAULT_USED_SETTING, c.DEFAULT_TEST_SET_NAME)
-        default_name = default_name_1 if default_name_1 else default_name_2
-        default_test_set = self.selection_test_sets.get(default_name, TestSetsConfig(name=default_name))
+        default_default_test_set = TestSetsConfig(name=c.DEFAULT_TEST_SET_NAME)
+        default_test_set = self.selection_test_sets.get(c.DEFAULT_TEST_SET_NAME, default_default_test_set)
         return default_test_set
     
-    def get_applicable_test_sets(self, dataset: str) -> list[str]:
-        applicable_test_sets = []
-        for test_set_name, test_set_config in self.selection_test_sets.items():
-            if test_set_config.datasets is None or dataset in test_set_config.datasets:
-                applicable_test_sets.append(test_set_name)
-        return applicable_test_sets
+    def get_default_configurables(self, dataset_name: str | None = None) -> dict[str, str]:
+        """
+        Return a dictionary of configurable name to its default value.
+        
+        If dataset_name is provided, merges project-level defaults with dataset-specific overrides.
+
+        Supports both list- and dict-shaped internal storage for configurables.
+        """
+        defaults: dict[str, str] = {}
+        for name, cfg in self.configurables.items():
+            defaults[name] = str(cfg.default)
+        
+        # Apply dataset-specific overrides if dataset_name is provided
+        if dataset_name is not None:
+            dataset_cfg = self.datasets.get(dataset_name)
+            if dataset_cfg:
+                for cfg_override in dataset_cfg.configurables:
+                    defaults[cfg_override.name] = cfg_override.default
+        
+        return defaults
 
 
 class ManifestIO:
 
     @classmethod
-    def load_from_file(cls, logger: _u.Logger, base_path: str, env_cfg: EnvironConfig) -> ManifestConfig:
+    def load_from_file(cls, logger: u.Logger, base_path: str, env_vars: dict[str, str]) -> ManifestConfig:
         start = time.time()
 
-        raw_content = _u.read_file(_u.Path(base_path, c.MANIFEST_FILE))
-        env_vars = env_cfg.get_all_env_vars()
-        content = _u.render_string(raw_content, base_path=base_path, env_vars=env_vars)
-        manifest_content = yaml.safe_load(content)
+        raw_content = u.read_file(u.Path(base_path, c.MANIFEST_FILE))
+        content = u.render_string(raw_content, base_path=base_path, env_vars=env_vars)
+        manifest_content: dict[str, Any] = yaml.safe_load(content)
+
+        auth_cfg: dict[str, Any] = manifest_content.get("authentication", {})
+        is_auth_required = bool(auth_cfg.get("is_required", False))
+
+        if is_auth_required:
+            # If authentication is required, assume PROTECTED when scope is not specified
+            # while explicitly forbidding PUBLIC (enforced in model validator)
+            datasets_raw = manifest_content.get("datasets", [])
+            for ds in datasets_raw:
+                if isinstance(ds, dict) and "scope" not in ds:
+                    ds["scope"] = "protected"
+        
         try:
-            manifest_cfg = ManifestConfig(base_path=base_path, env_cfg=env_cfg, **manifest_content)
+            manifest_cfg = ManifestConfig(base_path=base_path, **manifest_content)
         except ValidationError as e:
-            raise _u.ConfigurationError(f"Failed to process {c.MANIFEST_FILE} file. " + str(e)) from e
+            raise u.ConfigurationError(f"Failed to process {c.MANIFEST_FILE} file. " + str(e)) from e
         
         logger.log_activity_time(f"loading {c.MANIFEST_FILE} file", start)
         return manifest_cfg

squirrels/_model_builder.py

@@ -0,0 +1,69 @@
+from dataclasses import dataclass, field
+import duckdb, time
+
+from . import _utils as u, _connection_set as cs, _models as m
+
+
+@dataclass
+class ModelBuilder:
+    _datalake_db_path: str
+    _conn_set: cs.ConnectionSet
+    _static_models: dict[str, m.StaticModel]
+    _conn_args: cs.ConnectionsArgs = field(default_factory=lambda: cs.ConnectionsArgs(".", {}, {}))
+    _logger: u.Logger = field(default_factory=lambda: u.Logger(""))
+    
+    def _attach_connections(self, duckdb_conn: duckdb.DuckDBPyConnection) -> None:
+        for conn_name, conn_props in self._conn_set.get_connections_as_dict().items():
+            if not isinstance(conn_props, m.ConnectionProperties):
+                continue
+            attach_uri = conn_props.attach_uri_for_duckdb
+            if attach_uri is None:
+                continue # skip unsupported dialects
+            attach_stmt = f"ATTACH IF NOT EXISTS '{attach_uri}' AS db_{conn_name} (READ_ONLY)"
+            u.run_duckdb_stmt(self._logger, duckdb_conn, attach_stmt, redacted_values=[attach_uri])
+
+    async def _build_models(self, duckdb_conn: duckdb.DuckDBPyConnection, select: str | None, full_refresh: bool) -> None:
+        """
+        Compile and construct the build models as DuckDB tables.
+        """
+        # Compile the build models
+        models_list = self._static_models.values() if select is None else [self._static_models[select]]
+        for model in models_list:
+            model.compile_for_build(self._conn_args, self._static_models)
+
+        # Find all terminal nodes
+        terminal_nodes = set()
+        if select is None:
+            for model in models_list:
+                terminal_nodes.update(model.get_terminal_nodes_for_build(set()))
+            for model in models_list:
+                model.confirmed_no_cycles = False
+        else:
+            terminal_nodes.add(select)
+
+        # Run the build models
+        coroutines = []
+        for model_name in terminal_nodes:
+            model = self._static_models[model_name]
+            # await model.build_model(duckdb_conn, full_refresh)
+            coro = model.build_model(duckdb_conn, full_refresh)
+            coroutines.append(coro)
+        await u.asyncio_gather(coroutines)
+
+    async def build(self, full_refresh: bool, select: str | None) -> None:
+        start = time.time()
+
+        # Connect directly to DuckLake instead of attaching (supports concurrent connections)
+        duckdb_conn = u.create_duckdb_connection(self._datalake_db_path)
+        
+        try:
+            # Attach connections
+            self._attach_connections(duckdb_conn)
+
+            # Construct build models
+            await self._build_models(duckdb_conn, select, full_refresh)
+
+        finally:
+            duckdb_conn.close()
+
+        self._logger.log_activity_time("TOTAL TIME to build the Virtual Data Lake (VDL)", start)

squirrels/_model_configs.py

@@ -0,0 +1,74 @@
+from enum import Enum
+from pydantic import BaseModel, Field
+
+from . import _constants as c
+
+
+class ColumnCategory(Enum):
+    DIMENSION = "dimension"
+    MEASURE = "measure"
+    MISC = "misc"
+
+
+class ColumnConfig(BaseModel):
+    name: str = Field(description="The name of the column")
+    type: str = Field(default="", description="The type of the column such as 'string', 'integer', 'float', 'boolean', 'datetime', etc.")
+    condition: str = Field(default="", description="The condition of when the column is included")
+    description: str = Field(default="", description="The description of the column")
+    category: ColumnCategory = Field(default=ColumnCategory.MISC, description="The category of the column, either 'dimension', 'measure', or 'misc'")
+    depends_on: set[str] = Field(default_factory=set, description="List of dependent columns")
+    pass_through: bool = Field(default=False, description="Whether the column should be passed through to the federate")
+
+
+class ModelConfig(BaseModel):
+    description: str = Field(default="", description="The description of the model")
+    columns: list[ColumnConfig] = Field(default_factory=list, description="The columns of the model")
+
+
+class SeedConfig(ModelConfig):
+    cast_column_types: bool = Field(default=False, description="Whether the column types should be cast to the appropriate type")
+
+
+class ConnectionInterface(BaseModel):
+    connection: str | None = Field(default=None, description="The connection name of the source model / database view")
+
+    def finalize_connection(self, env_vars: dict[str, str]):
+        if self.connection is None:
+            self.connection = env_vars.get(c.SQRL_CONNECTIONS_DEFAULT_NAME_USED, "default")
+        return self
+
+    def get_connection(self) -> str:
+        assert self.connection is not None, "Connection must be set"
+        return self.connection
+    
+
+class QueryModelConfig(ModelConfig):
+    depends_on: set[str] = Field(default_factory=set, description="The dependencies of the model")
+
+
+class BuildModelConfig(QueryModelConfig):
+    materialization: str = Field(default="VIEW", description="The materialization of the model (ignored if Python model which is always a table)")
+
+    def get_sql_for_build(self, model_name: str, select_query: str) -> str:
+        if self.materialization.upper() == "TABLE":
+            materialization = "TABLE"
+        elif self.materialization.upper() == "VIEW":
+            materialization = "VIEW"
+        else:
+            raise ValueError(f"Invalid materialization: {self.materialization}")
+        
+        create_prefix = f"CREATE OR REPLACE {materialization} {model_name} AS\n\n"
+        return create_prefix + select_query
+
+
+class DbviewModelConfig(ConnectionInterface, QueryModelConfig):
+    translate_to_duckdb: bool = Field(default=False, description="Whether to translate the query to DuckDB and use DuckDB tables at runtime")
+
+
+class FederateModelConfig(QueryModelConfig):
+    eager: bool = Field(default=False, description="Whether the model should be materialized for SQL models")
+
+    def get_sql_for_create(self, model_name: str, select_query: str) -> str:
+        materialization = "TABLE" if self.eager else "VIEW"
+        create_prefix = f"CREATE {materialization} {model_name} AS\n\n"
+        return create_prefix + select_query

squirrels/_model_queries.py

@@ -0,0 +1,52 @@
+from abc import ABCMeta
+from dataclasses import dataclass, field
+from typing import Callable, Generic, TypeVar, Any
+import polars as pl, pandas as pd
+
+from ._arguments.run_time_args import BuildModelArgs
+from ._model_configs import ModelConfig
+
+
+# Input query file classes
+
+@dataclass(frozen=True)
+class QueryFile(metaclass=ABCMeta):
+    filepath: str
+    raw_query: Any
+
+@dataclass(frozen=True)
+class SqlQueryFile(QueryFile):
+    raw_query: str
+
+@dataclass(frozen=True)
+class PyQueryFile(QueryFile):
+    raw_query: Callable[[BuildModelArgs], pl.LazyFrame | pd.DataFrame]
+
+
+Q = TypeVar('Q', bound=QueryFile)
+M = TypeVar('M', bound=ModelConfig)
+
+@dataclass(frozen=True)
+class QueryFileWithConfig(Generic[Q, M]):
+    query_file: Q
+    config: M
+
+
+# Compiled query classes
+
+@dataclass
+class Query(metaclass=ABCMeta):
+    query: Any
+
+@dataclass
+class WorkInProgress(Query):
+    query: None = field(default=None, init=False)
+
+@dataclass
+class SqlModelQuery(Query):
+    query: str
+    is_duckdb: bool
+
+@dataclass
+class PyModelQuery(Query):
+    query: Callable[[], pl.LazyFrame | pd.DataFrame]