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

squirrels/_manifest.py

@@ -1,25 +1,58 @@
 from functools import cached_property
-from typing import Any
+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
 
 
-class ProjectVarsConfig(BaseModel, extra="allow"):
+class _ConfigWithNameBaseModel(BaseModel):
     name: str
+
+    @field_validator("name")
+    @classmethod
+    def validate_name(cls, v: str) -> str:
+        if not re.fullmatch(r"[A-Za-z0-9_-]+", v):
+            raise ValueError("Name must only contain alphanumeric characters, underscores, and dashes.")
+        return v
+
+
+class AuthType(Enum):
+    REQUIRED = "required"
+    OPTIONAL = "optional"
+    NOTSET = "notset"
+
+
+class AuthStrategy(Enum):
+    MANAGED = "managed"
+    EXTERNAL = "external"
+
+
+class ProjectVarsConfig(_ConfigWithNameBaseModel, extra="allow"):
+    major_version: int
     label: str = ""
     description: str = ""
-    major_version: int
+    auth_type: AuthType = AuthType.NOTSET
+    auth_strategy: AuthStrategy = AuthStrategy.MANAGED
+
+    @model_validator(mode="after")
+    def set_auth_strategy_defaults(self) -> Self:
+        if self.auth_strategy == AuthStrategy.EXTERNAL and self.auth_type == AuthType.OPTIONAL:
+            raise ValueError("auth_type can not be optional when auth_strategy is external")
+        
+        if self.auth_type == AuthType.NOTSET:
+            self.auth_type = AuthType.REQUIRED if self.auth_strategy == AuthStrategy.EXTERNAL else AuthType.OPTIONAL
+        
+        return self
 
     @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
 
 
@@ -35,14 +68,11 @@ class PackageConfig(BaseModel):
         return self
 
 
-class _ConfigWithNameBaseModel(BaseModel):
-    name: str
-
-
 class ConnectionTypeEnum(Enum):
     SQLALCHEMY = "sqlalchemy"
     CONNECTORX = "connectorx"
     ADBC = "adbc"
+    DUCKDB = "duckdb"
 
 
 class ConnectionProperties(BaseModel):
@@ -71,26 +101,32 @@ class ConnectionProperties(BaseModel):
 
     @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'] if dialect.lower().startswith(d)), None)
+        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.SQLALCHEMY:
+        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
-            sqlite_database = database if database is not None else ""
+            database_as_file = database if database is not None else ""
         else:
             url = urlparse(self.uri)
             host = url.hostname
@@ -98,22 +134,36 @@ class ConnectionProperties(BaseModel):
             username = url.username
             password = url.password
             database = url.path.lstrip('/')
-            sqlite_database = self.uri.replace(f"{self.dialect}://", "")
+            database_as_file = self.uri.replace(f"{self.dialect}://", "")
         
-        if self.dialect == 'sqlite':
-            return sqlite_database
-        elif self.dialect in ('postgres', 'mysql'):
-            return f"dbname={database} user={username} password={password} host={host} port={port}"
+        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:
-            return None
+            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)
+    def finalize_uri(self, project_path: str) -> Self:
+        self.uri = self.uri.format(project_path=project_path)
         return self
 
 
+class ConfigurableOverride(BaseModel):
+    name: str
+    default: str
+
+
+class ConfigurablesConfig(ConfigurableOverride):
+    label: str = ""
+    description: str = ""
+
+
 class ParametersConfig(BaseModel):
     type: str
     factory: str
@@ -129,8 +179,10 @@ class PermissionScope(Enum):
 class AnalyticsOutputConfig(_ConfigWithNameBaseModel):
     label: str = ""
     description: str = ""
-    scope: PermissionScope = PermissionScope.PUBLIC
+    scope: PermissionScope | None = None
     parameters: list[str] | None = Field(default=None, description="The list of parameter names used by the dataset/dashboard")
+    configurables: list[ConfigurableOverride] = Field(default_factory=list)
+    project_configurables: dict[str, Any] | None = Field(default=None, exclude=True)
 
     @model_validator(mode="after")
     def finalize_label(self) -> Self:
@@ -140,7 +192,9 @@ class AnalyticsOutputConfig(_ConfigWithNameBaseModel):
 
     @field_validator("scope", mode="before")
     @classmethod
-    def validate_scope(cls, value: str, info: ValidationInfo) -> PermissionScope:
+    def validate_scope(cls, value: Any, info: ValidationInfo) -> PermissionScope | None:
+        if value is None:
+            return None
         try:
             return PermissionScope[str(value).upper()]
         except KeyError as e:
@@ -148,15 +202,23 @@ class AnalyticsOutputConfig(_ConfigWithNameBaseModel):
             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 DatasetTraitConfig(_ConfigWithNameBaseModel):
-    default: Any 
+    @model_validator(mode="after")
+    def validate_configurables(self) -> Self:
+        if self.project_configurables is not None:
+            for cfg_override in self.configurables:
+                if cfg_override.name not in self.project_configurables:
+                    # Determine if it's a dataset or dashboard for better error message
+                    class_name = self.__class__.__name__
+                    type_str = "Dataset" if "Dataset" in class_name else "Dashboard" if "Dashboard" in class_name else "Asset"
+                    raise ValueError(
+                        f'{type_str} "{self.name}" references configurable "{cfg_override.name}" which is not defined '
+                        f'in the project configurables'
+                    )
+        return self
 
 
 class DatasetConfig(AnalyticsOutputConfig):
     model: str = ""
-    traits: dict = Field(default_factory=dict)
-    default_test_set: str = ""
     
     def __hash__(self) -> int:
         return hash("dataset_"+self.name)
@@ -168,14 +230,14 @@ class DatasetConfig(AnalyticsOutputConfig):
         return self
 
 
+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
-    user_attributes: dict[str, Any] | None = None
+    user: TestSetsUserConfig = Field(default_factory=TestSetsUserConfig)
     parameters: dict[str, Any] = Field(default_factory=dict)
-
-    @property
-    def is_authenticated(self) -> bool:
-        return self.user_attributes is not None
+    configurables: dict[str, Any] = Field(default_factory=dict)
 
 
 class ManifestConfig(BaseModel):
@@ -183,11 +245,10 @@ class ManifestConfig(BaseModel):
     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)
-    dataset_traits: dict[str, DatasetTraitConfig] = Field(default_factory=dict)
     datasets: dict[str, DatasetConfig] = Field(default_factory=dict)
-    base_path: str = "."
-    env_vars: dict[str, str] = Field(default_factory=dict)
+    project_path: str = "."
 
     @field_validator("packages")
     @classmethod
@@ -199,13 +260,13 @@ class ManifestConfig(BaseModel):
             set_of_directories.add(package.directory)
         return packages
     
-    @field_validator("connections", "selection_test_sets", "dataset_traits", "datasets", 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
@@ -216,62 +277,79 @@ class ManifestConfig(BaseModel):
     @model_validator(mode="after")
     def finalize_connections(self) -> Self:
         for conn in self.connections.values():
-            conn.finalize_uri(self.base_path)
+            conn.finalize_uri(self.project_path)
         return self
     
     @model_validator(mode="after")
-    def validate_dataset_traits(self) -> Self:
-        for dataset_name, dataset in self.datasets.items():
-            # Validate that all trait keys in dataset.traits exist in dataset_traits
-            for trait_key in dataset.traits.keys():
-                if trait_key not in self.dataset_traits:
-                    raise ValueError(
-                        f'Dataset "{dataset_name}" references undefined trait "{trait_key}". '
-                        f'Traits must be defined with a default value in the dataset_traits section.'
-                    )
-            
-            # Set default values for any traits that are missing
-            for trait_name, trait_config in self.dataset_traits.items():
-                if trait_name not in dataset.traits:
-                    dataset.traits[trait_name] = trait_config.default
-                    
+    def validate_authentication_and_scopes(self) -> Self:
+        """
+        Enforce authentication rules:
+        - Set default scope based on auth_type if not specified.
+        - If auth_type is REQUIRED, no dataset may be PUBLIC.
+        """
+        is_auth_required = self.project_variables.auth_type == AuthType.REQUIRED
+        default_scope = PermissionScope.PROTECTED if is_auth_required else PermissionScope.PUBLIC
+        
+        for ds in self.datasets.values():
+            if ds.scope is None:
+                ds.scope = default_scope
+
+        if is_auth_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.\n  "
+                    f"Update the scope for the following datasets: {invalid}\n "
+                )
         return self
     
-    def get_default_test_set(self, dataset_name: str) -> TestSetsConfig:
+    @model_validator(mode="after")
+    def validate_dataset_configurables(self) -> Self:
+        """
+        Validate that dataset configurables reference valid project-level configurables.
+        """
+        for dataset_cfg in self.datasets.values():
+            dataset_cfg.project_configurables = self.configurables
+            dataset_cfg.validate_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.env_vars.get(c.SQRL_TEST_SETS_DEFAULT_NAME_USED, "default")
-        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_traits(self) -> dict[str, Any]:
-        default_traits = {}
-        for trait_name, trait_config in self.dataset_traits.items():
-            default_traits[trait_name] = trait_config.default
-        return default_traits
+    def get_default_configurables(self, *, overrides: list[ConfigurableOverride] = []) -> dict[str, str]:
+        """
+        Return a dictionary of configurable name to its default value.
+        
+        Arguments:
+            overrides: A list of ConfigurableOverride objects to merge with the project-level defaults.
+        """
+        defaults: dict[str, str] = {}
+        for name, cfg in self.configurables.items():
+            defaults[name] = str(cfg.default)
+        
+        # Apply explicit overrides if provided
+        for cfg_override in overrides:
+            defaults[cfg_override.name] = cfg_override.default
+        
+        return defaults
 
 
 class ManifestIO:
-
     @classmethod
-    def load_from_file(cls, logger: u.Logger, base_path: str, env_vars: dict[str, str]) -> ManifestConfig:
+    def load_from_file(cls, logger: u.Logger, project_path: str, env_vars_unformatted: dict[str, str]) -> ManifestConfig:
         start = time.time()
 
-        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 = yaml.safe_load(content)
+        raw_content = u.read_file(u.Path(project_path, c.MANIFEST_FILE))
+        content = u.render_string(raw_content, project_path=project_path, env_vars=env_vars_unformatted)
+        manifest_content: dict[str, Any] = yaml.safe_load(content)
+
         try:
-            manifest_cfg = ManifestConfig(base_path=base_path, **manifest_content)
+            manifest_cfg = ManifestConfig(project_path=project_path, **manifest_content)
         except ValidationError as e:
             raise u.ConfigurationError(f"Failed to process {c.MANIFEST_FILE} file. " + str(e)) from e