squirrels 0.1.0__py3-none-any.whl → 0.6.0.post0__py3-none-any.whl

squirrels/_manifest.py

@@ -1,187 +1,357 @@
-from typing import List, Dict, Any, Optional, Union
-from pathlib import Path
-from sqlalchemy import Engine, create_engine
-import yaml
-
-from squirrels import _constants as c, _utils
-from squirrels._credentials_manager import Credential, squirrels_config_io
-from squirrels._utils import ConfigurationError, InvalidInputError
-
-
-class Manifest:
-    def __init__(self, parms: Dict) -> None:
-        self._parms = parms
-    
-    @classmethod
-    def from_yaml_str(cls, parms_str: str):
-        parms = yaml.safe_load(parms_str)
-        return cls(parms)
-
-    @classmethod
-    def from_file(cls, manifest_path: str):
-        with open(manifest_path, 'r') as f:
-            parms_str = f.read()
-        
-        return Manifest.from_yaml_str(parms_str)
-    
-    def get_proj_vars(self) -> Dict[str, Any]:
-        return self._parms.get(c.PROJ_VARS_KEY, dict())
-    
-    def get_modules(self) -> List[str]:
-        return self._parms.get(c.MODULES_KEY, list())
-    
-    def _get_required_field(self, key: str) -> Any:
-        try:
-            return self._parms[key]
-        except KeyError as e:
-            raise ConfigurationError(f'Field "{key}" not found in squirrels.yaml') from e
-    
-    def get_base_path(self) -> str:
-        project_vars = self.get_proj_vars()
-        try:
-            product = project_vars[c.PRODUCT_KEY]
-            major_version = project_vars[c.MAJOR_VERSION_KEY]
-        except KeyError as e:
-            raise ConfigurationError("Could not construct API endpoint as 'product' and 'major_version'" + 
-                "were not specified in project variables") from e
-        base_path = f"/{product}/v{major_version}"
-        return base_path
-    
-    def get_db_connections(self, test_creds: Dict[str, Credential] = None) -> Dict[str, Engine]:
-        configs: Dict[str, Dict[str, str]] = self._parms.get(c.DB_CONNECTIONS_KEY, {})
-        output = {}
-        for key, config in configs.items():
-            cred_key = config.get(c.DB_CREDENTIALS_KEY)
-            if cred_key is None:
-                cred = Credential("", "")
-            elif test_creds is not None:
-                cred = test_creds[cred_key]
-            else:
-                cred = squirrels_config_io.get_credential(cred_key)
-            url = config[c.URL_KEY].replace("${username}", cred.username).replace("${password}", cred.password)
-            output[key] = create_engine(url)
-        return output
-
-    def _get_dataset_parms(self, dataset: str) -> Dict[str, Any]:
-        try:
-            return self._get_required_field(c.DATASETS_KEY)[dataset]
-        except KeyError as e:
-            raise InvalidInputError(f'No such dataset named "{dataset}" exists') from e
-        
-    def _get_required_field_from_dataset_parms(self, dataset: str, key: str):
-        try:
-            return self._get_dataset_parms(dataset)[key]
-        except KeyError as e:
-            raise ConfigurationError(f'The "{key}" field is not defined for dataset "{dataset}"') from e
-    
-    def _get_all_database_view_parms(self, dataset: str) -> Dict[str, Dict[str, str]]:
-        return self._get_required_field_from_dataset_parms(dataset, c.DATABASE_VIEWS_KEY)
-    
-    def get_all_dataset_names(self) -> str:
-        datasets: Dict[str, Any] = self._get_required_field(c.DATASETS_KEY)
-        return list(datasets.keys())
-    
-    def get_dataset_folder(self, dataset: str) -> Path:
-        return _utils.join_paths(c.DATASETS_FOLDER, dataset)
-        
-    def get_dataset_args(self, dataset: str) -> Dict[str, Any]:
-        dataset_args = self._get_dataset_parms(dataset).get("args", {})
-        full_args = {**self.get_proj_vars(), **dataset_args}
-        return full_args
-    
-    def get_all_database_view_names(self, dataset: str) -> List[str]:
-        all_database_views = self._get_all_database_view_parms(dataset)
-        return list(all_database_views.keys())
-    
-    def get_database_view_file(self, dataset: str, database_view: str) -> Path:
-        database_view_parms = self._get_all_database_view_parms(dataset)[database_view]
-        if isinstance(database_view_parms, str):
-            db_view_file = database_view_parms
-        else:
-            try:
-                db_view_file = database_view_parms[c.FILE_KEY]
-            except KeyError as e:
-                raise ConfigurationError(f'The "{c.FILE_KEY}" field is not defined for "{database_view}" in dataset "{dataset}"') from e
-        dataset_folder = self.get_dataset_folder(dataset)
-        return _utils.join_paths(dataset_folder, db_view_file)
-    
-    def get_view_args(self, dataset: str, database_view: str = None) -> Dict[str, Any]:
-        dataset_args = self.get_dataset_args(dataset)
-        if database_view is None:
-            view_parms: Dict[str, Any] = self._get_required_field_from_dataset_parms(dataset, c.FINAL_VIEW_KEY)
-        else:
-            view_parms: Dict[str, Any] = self._get_all_database_view_parms(dataset)[database_view]
-        view_args: Dict[str, Any] = {} if isinstance(view_parms, str) else view_parms.get("args", {})
-        full_args = {**dataset_args, **view_args}
-        return full_args
-
-    def get_database_view_db_connection(self, dataset: str, database_view: str) -> Optional[str]:
-        database_view_parms = self._get_all_database_view_parms(dataset)[database_view]
-        if isinstance(database_view_parms, str):
-            db_connection = c.DEFAULT_DB_CONN 
-        else: 
-            db_connection = database_view_parms.get(c.DB_CONNECTION_KEY, c.DEFAULT_DB_CONN)
-        return db_connection
-    
-    def get_dataset_label(self, dataset: str) -> str:
-        return self._get_required_field_from_dataset_parms(dataset, c.DATASET_LABEL_KEY)
-    
-    def get_dataset_final_view_file(self, dataset: str) -> Union[str, Path]:
-        final_view_parms: Dict[str, Any] = self._get_required_field_from_dataset_parms(dataset, c.FINAL_VIEW_KEY)
-        if isinstance(final_view_parms, str):
-            final_view_file = final_view_parms
-        else:
-            try:
-                final_view_file = final_view_parms[c.FILE_KEY]
-            except KeyError as e:
-                raise ConfigurationError(f'The "{c.FILE_KEY}" field is not defined for the final view') from e
-        
-        database_views = self.get_all_database_view_names(dataset)
-        if final_view_file in database_views:
-            return final_view_file
-        else:
-            dataset_path = self.get_dataset_folder(dataset)
-            return _utils.join_paths(dataset_path, final_view_file)
-
-    def get_setting(self, key: str, default: Any) -> Any:
-        settings: Dict[str, Any] = self._parms.get(c.SETTINGS_KEY, dict())
-        return settings.get(key, default)
-    
-    def get_catalog(self, parameters_path: str, results_path: str) -> Any:
-        """
-        Gets the component of the catalog API response that's generated by this manifest
-
-        Parameters:
-            parameters_path: The path to the parameters API endpoint
-            results_path: The path to the results API endpoint
-        
-        Returns:
-            A JSON response for the catalog API
-        """
-        datasets_info = []
-        for dataset in self.get_all_dataset_names():
-            dataset_normalized = _utils.normalize_name_for_api(dataset)
-            datasets_info.append({
-                'name': dataset,
-                'label': self.get_dataset_label(dataset),
-                'parameters_path': parameters_path.format(dataset=dataset_normalized),
-                'result_path': results_path.format(dataset=dataset_normalized),
-                'first_minor_version': 0
-            })
-        
-        project_vars = self.get_proj_vars()
-        return {
-            'response_version': 0,
-            'products': [{
-                'name': project_vars[c.PRODUCT_KEY],
-                'versions': [{
-                    'major_version': project_vars[c.MAJOR_VERSION_KEY],
-                    'latest_minor_version': project_vars[c.MINOR_VERSION_KEY],
-                    'datasets': datasets_info
-                }]
-            }]
-        }
-
-
-def _from_file():
-    return Manifest.from_file(c.MANIFEST_FILE)
+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, re
+
+from . import _constants as c, _utils as u
+
+
+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 = ""
+    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 = u.to_title_case(self.name)
+        return self
+
+
+class PackageConfig(BaseModel):
+    git: str
+    revision: str
+    directory: str = ""
+
+    @model_validator(mode="after")
+    def finalize_directory(self) -> Self:
+        if self.directory == "":
+            self.directory = self.git.split('/')[-1].removesuffix('.git')
+        return self
+
+
+class ConnectionTypeEnum(Enum):
+    SQLALCHEMY = "sqlalchemy"
+    CONNECTORX = "connectorx"
+    ADBC = "adbc"
+    DUCKDB = "duckdb"
+
+
+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, 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
+    arguments: dict[str, Any]
+
+
+class PermissionScope(Enum):
+    PUBLIC = 0
+    PROTECTED = 1
+    PRIVATE = 2
+
+
+class AnalyticsOutputConfig(_ConfigWithNameBaseModel):
+    label: str = ""
+    description: str = ""
+    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:
+        if self.label == "":
+            self.label = self.name
+        return self
+
+    @field_validator("scope", mode="before")
+    @classmethod
+    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:
+            name = info.data.get("name")
+            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
+
+    @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 = ""
+    
+    def __hash__(self) -> int:
+        return hash("dataset_"+self.name)
+
+    @model_validator(mode="after")
+    def finalize_model(self) -> Self:
+        if self.model == "":
+            self.model = self.name
+        return self
+
+
+class TestSetsUserConfig(BaseModel):
+    access_level: Literal["admin", "member", "guest"] = "guest"
+    custom_fields: dict[str, Any] = Field(default_factory=dict)
+
+class TestSetsConfig(_ConfigWithNameBaseModel):
+    user: TestSetsUserConfig = Field(default_factory=TestSetsUserConfig)
+    parameters: dict[str, Any] = Field(default_factory=dict)
+    configurables: dict[str, Any] = Field(default_factory=dict)
+
+
+class ManifestConfig(BaseModel):
+    project_variables: ProjectVarsConfig
+    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)
+    datasets: dict[str, DatasetConfig] = Field(default_factory=dict)
+    project_path: str = "."
+
+    @field_validator("packages")
+    @classmethod
+    def package_directories_are_unique(cls, packages: list[PackageConfig]) -> list[PackageConfig]:
+        set_of_directories = set()
+        for package in packages:
+            if package.directory in set_of_directories:
+                raise ValueError(f'In the packages section, multiple target directories found for "{package.directory}"')
+            set_of_directories.add(package.directory)
+        return packages
+    
+    @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 = 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
+        else:
+            values_as_dict = values
+        return values_as_dict
+    
+    @model_validator(mode="after")
+    def finalize_connections(self) -> Self:
+        for conn in self.connections.values():
+            conn.finalize_uri(self.project_path)
+        return self
+    
+    @model_validator(mode="after")
+    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
+    
+    @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_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_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, project_path: str, env_vars_unformatted: dict[str, str]) -> ManifestConfig:
+        start = time.time()
+
+        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(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
+        
+        logger.log_activity_time(f"loading {c.MANIFEST_FILE} file", start)
+        return manifest_cfg