squirrels 0.3.2__py3-none-any.whl → 0.4.0__py3-none-any.whl

squirrels/_manifest.py

@@ -1,141 +1,63 @@
-from typing import Optional
-from dataclasses import dataclass, field
+from typing import Any
+from typing_extensions import Self
 from enum import Enum
-import yaml
+from pydantic import BaseModel, Field, field_validator, model_validator, ValidationInfo, ValidationError
+import yaml, time
 
-from . import _constants as c, _utils as u
-from ._environcfg import EnvironConfigIO
-from ._timer import timer, time
+from . import _constants as c, _utils as _u
+from ._environcfg import EnvironConfig
 
 
-@dataclass
-class ManifestComponentConfig:
-    @classmethod
-    def _validate_required(cls, data: dict, required_keys: list[str], section: str):
-        for key in required_keys:
-            if key not in data:
-                raise u.ConfigurationError(f'In {c.MANIFEST_FILE}, required field missing in {section}: {key}')
-    
-    @classmethod
-    def from_dict(cls, kwargs: dict):
-        return cls()
+class ProjectVarsConfig(BaseModel, extra="allow"):
+    name: str
+    label: str = ""
+    major_version: int
 
+    @model_validator(mode="after")
+    def finalize_label(self) -> Self:
+        if self.label == "":
+            self.label = self.name
+        return self
 
-@dataclass
-class ProjectVarsConfig(ManifestComponentConfig):
-    data: dict
 
-    def __post_init__(self):
-        required_keys = [c.PROJECT_NAME_KEY, c.MAJOR_VERSION_KEY]
-        self._validate_required(self.data, required_keys, c.PROJ_VARS_KEY)
-        
-        integer_keys = [c.MAJOR_VERSION_KEY]
-        for key in integer_keys:
-            if key in self.data and not isinstance(self.data[key], int):
-                raise u.ConfigurationError(f'Project variable "{key}" must be an integer')
-    
-    @classmethod
-    def from_dict(cls, kwargs: dict):
-        return cls(kwargs)
-    
-    def get_name(self) -> str:
-        return str(self.data[c.PROJECT_NAME_KEY])
-    
-    def get_label(self) -> str:
-        return str(self.data.get(c.PROJECT_LABEL_KEY, self.get_name()))
-    
-    def get_major_version(self) -> int:
-        return self.data[c.MAJOR_VERSION_KEY]
-
-
-@dataclass
-class PackageConfig(ManifestComponentConfig):
-    git_url: str
-    directory: str
+class PackageConfig(BaseModel):
+    git: str
     revision: str
+    directory: str = ""
 
-    @classmethod
-    def from_dict(cls, kwargs: dict):
-        cls._validate_required(kwargs, [c.PACKAGE_GIT_KEY, c.PACKAGE_REVISION_KEY], c.PACKAGES_KEY)
-        git_url = str(kwargs[c.PACKAGE_GIT_KEY])
-        directory_raw = kwargs.get(c.PACKAGE_DIRECTORY_KEY)
-        directory = git_url.split('/')[-1].removesuffix('.git') if directory_raw is None else str(directory_raw)
-        revision = str(kwargs[c.PACKAGE_REVISION_KEY])
-        return cls(git_url, directory, revision)
+    @model_validator(mode="after")
+    def finalize_directory(self) -> Self:
+        if self.directory == "":
+            self.directory = self.git.split('/')[-1].removesuffix('.git')
+        return self
 
 
-@dataclass
-class DbConnConfig(ManifestComponentConfig):
+class _ConfigWithNameBaseModel(BaseModel):
     name: str
+
+
+class DbConnConfig(_ConfigWithNameBaseModel):
+    credential: str | None = None
     url: str
 
-    @classmethod
-    def from_dict(cls, kwargs: dict):
-        cls._validate_required(kwargs, [c.DB_CONN_NAME_KEY, c.DB_CONN_URL_KEY], c.DB_CONNECTIONS_KEY)
-        name = str(kwargs[c.DB_CONN_NAME_KEY])
-        credential_key = kwargs.get(c.DB_CONN_CRED_KEY)
-        username, password = EnvironConfigIO.obj.get_credential(credential_key)
-        url = str(kwargs[c.DB_CONN_URL_KEY]).format(username=username, password=password)
-        return cls(name, url)
+    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)
+        return self
 
 
-@dataclass
-class ParametersConfig(ManifestComponentConfig):
+class ParametersConfig(BaseModel):
     type: str
     factory: str
-    arguments: dict
-
-    @classmethod
-    def from_dict(cls, kwargs: dict):
-        all_keys = [c.PARAMETER_TYPE_KEY, c.PARAMETER_FACTORY_KEY, c.PARAMETER_ARGS_KEY]
-        cls._validate_required(kwargs, all_keys, c.PARAMETERS_KEY)
-        args = {key: kwargs[key] for key in all_keys}
-        return cls(**args)
-
-
-@dataclass
-class TestSetsConfig(ManifestComponentConfig):
-    name: str
-    datasets: Optional[list[str]] = None
-    is_authenticated: bool = False
-    user_attributes: dict = field(default_factory=dict)
-    parameters: dict = field(default_factory=dict)
+    arguments: dict[str, Any]
 
-    @classmethod
-    def from_dict(cls, kwargs: dict):
-        cls._validate_required(kwargs, [c.TEST_SET_NAME_KEY], c.TEST_SETS_KEY)
-        name = str(kwargs[c.TEST_SET_NAME_KEY])
-        datasets = kwargs.get(c.TEST_SET_DATASETS_KEY)
-        is_authenticated = (c.TEST_SET_USER_ATTR_KEY in kwargs)
-        user_attributes = kwargs.get(c.TEST_SET_USER_ATTR_KEY, {})
-        parameters = kwargs.get(c.TEST_SET_PARAMETERS_KEY, {})
-        return cls(name, datasets, is_authenticated, user_attributes, parameters)
-
-
-@dataclass
-class DbviewConfig(ManifestComponentConfig):
-    name: str
-    connection_name: Optional[str]
 
-    @classmethod
-    def from_dict(cls, kwargs: dict):
-        cls._validate_required(kwargs, [c.DBVIEW_NAME_KEY], c.DBVIEWS_KEY)
-        name = str(kwargs[c.DBVIEW_NAME_KEY])
-        connection_name = str(kwargs.get(c.DBVIEW_CONN_KEY))
-        return cls(name, connection_name)
+class DbviewConfig(_ConfigWithNameBaseModel):
+    connection_name: str | None = None
 
 
-@dataclass
-class FederateConfig(ManifestComponentConfig):
-    name: str
-    materialized: Optional[str]
-
-    @classmethod
-    def from_dict(cls, kwargs: dict):
-        cls._validate_required(kwargs, [c.FEDERATE_NAME_KEY], c.FEDERATES_KEY)
-        name = str(kwargs[c.FEDERATE_NAME_KEY])
-        materialized = str(kwargs.get(c.MATERIALIZED_KEY))
-        return cls(name, materialized)
+class FederateConfig(_ConfigWithNameBaseModel):
+    materialized: str | None = None
 
 
 class DatasetScope(Enum):
@@ -143,102 +65,153 @@ class DatasetScope(Enum):
     PROTECTED = 1
     PRIVATE = 2
 
-@dataclass
-class DatasetsConfig(ManifestComponentConfig):
-    name: str
-    label: str
-    model: str
-    scope: DatasetScope
-    parameters: list[str]
-    traits: dict
-    default_test_set: Optional[str]
 
+class AnalyticsOutputConfig(_ConfigWithNameBaseModel):
+    label: str = ""
+    description: str = ""
+    scope: DatasetScope = DatasetScope.PUBLIC
+    parameters: list[str] = Field(default_factory=list)
+
+    @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 from_dict(cls, kwargs: dict):
-        cls._validate_required(kwargs, [c.DATASET_NAME_KEY], c.DATASETS_KEY)
-        name = str(kwargs[c.DATASET_NAME_KEY])
-        label = str(kwargs.get(c.DATASET_LABEL_KEY, name))
-        model = str(kwargs.get(c.DATASET_MODEL_KEY, name))
-        scope_raw = kwargs.get(c.DATASET_SCOPE_KEY)
+    def validate_scope(cls, value: str, info: ValidationInfo) -> DatasetScope:
         try:
-            scope = DatasetScope[str(scope_raw).upper()] if scope_raw is not None else DatasetScope.PUBLIC
+            return DatasetScope[str(value).upper()]
         except KeyError as e:
+            name = info.data.get("name")
             scope_list = [scope.name.lower() for scope in DatasetScope]
-            raise u.ConfigurationError(f'Scope not found for dataset "{name}". Scope must be one of {scope_list}') from e
-        
-        parameters = kwargs.get(c.DATASET_PARAMETERS_KEY, [])
-        traits = kwargs.get(c.DATASET_TRAITS_KEY, {})
-        default_test_set = kwargs.get(c.DATASET_DEFAULT_TEST_SET_KEY)
-        return cls(name, label, model, scope, parameters, traits, default_test_set)
+            raise ValueError(f'Scope "{value}" is invalid for dataset/dashboard "{name}". Scope must be one of {scope_list}') from e
 
 
-@dataclass
-class _ManifestConfig:
-    project_variables: ProjectVarsConfig
-    packages: list[PackageConfig]
-    connections: dict[str, DbConnConfig]
-    parameters: list[ParametersConfig]
-    selection_test_sets: dict[str, TestSetsConfig]
-    dbviews: dict[str, DbviewConfig]
-    federates: dict[str, FederateConfig]
-    datasets: dict[str, DatasetsConfig]
-    settings: dict
+class DatasetConfig(AnalyticsOutputConfig):
+    model: str = ""
+    traits: dict = Field(default_factory=dict)
+    default_test_set: str = ""
+    
+    def __hash__(self) -> int:
+        return hash("dataset_"+self.name)
 
-    @classmethod
-    def _create_configs_as_dict(cls, config_cls: ManifestComponentConfig, kwargs: dict, section_key: str, name_key: str) -> dict:
-        configs_dict = {}
-        for x in kwargs.get(section_key, []):
-            name = x[name_key]
-            if name in configs_dict:
-                raise u.ConfigurationError(f'In the "{section_key}" section of {c.MANIFEST_FILE}, the name/identifier "{name}" was specified multiple times')
-            configs_dict[name] = config_cls.from_dict(x)
-        return configs_dict
+    @model_validator(mode="after")
+    def finalize_model(self) -> Self:
+        if self.model == "":
+            self.model = self.name
+        return self
 
-    @classmethod
-    def from_dict(cls, kwargs: dict):
-        settings: dict = kwargs.get(c.SETTINGS_KEY, {})
 
-        try:
-            proj_vars = ProjectVarsConfig(kwargs[c.PROJ_VARS_KEY])
-        except KeyError as e:
-            raise u.ConfigurationError(f'In {c.MANIFEST_FILE}, section for {c.PROJ_VARS_KEY} is required') from e
-        
-        packages = [PackageConfig.from_dict(x) for x in kwargs.get(c.PACKAGES_KEY, [])]
-        all_package_dirs = set()
-        for package in packages:
-            if package.directory in all_package_dirs:
-                raise u.ConfigurationError(f'In the "{c.PACKAGES_KEY}" section of {c.MANIFEST_FILE}, multiple target directories found for "{package.directory}"')
-            all_package_dirs.add(package.directory)
+class DashboardConfig(AnalyticsOutputConfig):
+    def __hash__(self) -> int:
+        return hash("dashboard_"+self.name)
 
-        db_conns = cls._create_configs_as_dict(DbConnConfig, kwargs, c.DB_CONNECTIONS_KEY, c.DB_CONN_NAME_KEY)
-        params = [ParametersConfig.from_dict(x) for x in kwargs.get(c.PARAMETERS_KEY, [])]
 
-        test_sets = cls._create_configs_as_dict(TestSetsConfig, kwargs, c.TEST_SETS_KEY, c.TEST_SET_NAME_KEY)
-        dbviews = cls._create_configs_as_dict(DbviewConfig, kwargs, c.DBVIEWS_KEY, c.DBVIEW_NAME_KEY)
-        federates = cls._create_configs_as_dict(FederateConfig, kwargs, c.FEDERATES_KEY, c.FEDERATE_NAME_KEY)
-        datasets = cls._create_configs_as_dict(DatasetsConfig, kwargs, c.DATASETS_KEY, c.DATASET_NAME_KEY)
+class TestSetsConfig(_ConfigWithNameBaseModel):
+    datasets: list[str] | None = None
+    is_authenticated: bool = False
+    user_attributes: dict[str, Any] = Field(default_factory=dict)
+    parameters: dict[str, Any] = Field(default_factory=dict)
 
-        return cls(proj_vars, packages, db_conns, params, test_sets, dbviews, federates, datasets, settings)
+    @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 get_default_test_set(self, dataset_name: str) -> tuple[str, dict]:
-        default_1 = self.datasets[dataset_name].default_test_set
-        default_2 = self.settings.get(c.TEST_SET_DEFAULT_USED_SETTING, c.DEFAULT_TEST_SET_NAME)
-        default_name = default_1 if default_1 is not None else default_2
-        default_test_set = self.selection_test_sets.get(default_name, TestSetsConfig.from_dict({c.TEST_SET_NAME_KEY: default_name}))
-        return default_name, default_test_set
+    def do_use_duckdb(self) -> bool:
+        return self.data.get(c.IN_MEMORY_DB_SETTING, c.SQLITE) == c.DUCKDB
+
+
+class ManifestConfig(BaseModel):
+    env_cfg: EnvironConfig
+    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)
+    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 = "."
+
+    @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", "dbviews", "federates", "datasets", "dashboards", 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"]
+                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_url(self.base_path, self.env_cfg)
+        return self
+    
+    @property
+    def settings_obj(self) -> Settings:
+        return Settings(data=self.settings)
+
+    def get_default_test_set(self, dataset_name: str) -> 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))
+        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
 
 
 class ManifestIO:
-    obj: _ManifestConfig
 
     @classmethod
-    def LoadFromFile(cls) -> None:
-        EnvironConfigIO.LoadFromFile()
-        
+    def load_from_file(cls, logger: _u.Logger, base_path: str, env_cfg: EnvironConfig) -> ManifestConfig:
         start = time.time()
-        raw_content = u.read_file(c.MANIFEST_FILE)
-        env_vars = EnvironConfigIO.obj.get_all_env_vars()
-        content = u.render_string(raw_content, env_vars=env_vars, **env_vars) # TODO: deprecate **env_vars
-        proj_config = yaml.safe_load(content)
-        cls.obj = _ManifestConfig.from_dict(proj_config)
-        timer.add_activity_time(f"loading {c.MANIFEST_FILE} file", start)
+
+        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)
+        try:
+            manifest_cfg = ManifestConfig(base_path=base_path, env_cfg=env_cfg, **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