diff --git a/cognite/neat/_graph/transformers/__init__.py b/cognite/neat/_graph/transformers/__init__.py index 58ff69783..317c3e7b1 100644 --- a/cognite/neat/_graph/transformers/__init__.py +++ b/cognite/neat/_graph/transformers/__init__.py @@ -1,4 +1,4 @@ -from ._base import BaseTransformerStandardised +from ._base import BaseTransformer, BaseTransformerStandardised from ._classic_cdf import ( AddAssetDepth, AssetEventConnector, @@ -27,6 +27,7 @@ "AssetSequenceConnector", "AssetTimeSeriesConnector", "AttachPropertyFromTargetToSource", + "BaseTransformer", "ConnectionToLiteral", "ConvertLiteral", "LiteralToEntity", diff --git a/cognite/neat/_session/_read.py b/cognite/neat/_session/_read.py index 4c2cb838e..88a173985 100644 --- a/cognite/neat/_session/_read.py +++ b/cognite/neat/_session/_read.py @@ -705,12 +705,6 @@ def pump_example(self) -> IssueList: def core_data_model(self) -> IssueList: """Reads the core data model example into the NeatSession.""" - - self._state._raise_exception_if_condition_not_met( - "Read Core Data Model example", - empty_rules_store_required=True, - ) - cdm_v1 = DataModelId.load(("cdf_cdm", "CogniteCore", "v1")) importer: importers.DMSImporter = importers.DMSImporter.from_data_model_id(self._get_client, cdm_v1) - return self._state.rule_import(importer) + return self._state.change(importer, "Read Core Data Model example") diff --git a/cognite/neat/_session/_state.py b/cognite/neat/_session/_state.py index 361cddb57..af1d2712c 100644 --- a/cognite/neat/_session/_state.py +++ b/cognite/neat/_session/_state.py @@ -11,7 +11,8 @@ from cognite.neat._rules.transformers import ( VerifiedRulesTransformer, ) -from cognite.neat._store import NeatGraphStore, NeatRulesStore +from cognite.neat._store import NeatGraphStore, NeatRulesStore, NeatStoreManager +from cognite.neat._store._manager import Action from cognite.neat._utils.upload import UploadResultList from .exceptions import NeatSessionError, _session_method_wrapper @@ -24,12 +25,19 @@ def __init__( storage_path: Path | None = None, client: NeatClient | None = None, ) -> None: - self.instances = InstancesState(store_type, storage_path=storage_path) self.rule_store = NeatRulesStore() - self.last_reference: DMSRules | InformationRules | None = None + self.instances = InstancesState(store_type, storage_path=storage_path) + self.manager = NeatStoreManager(self.instances.store, self.rule_store) self.client = client + self.last_reference: DMSRules | InformationRules | None = None self.quoted_source_identifiers = False + def change(self, action: Action, description: str | None = None) -> IssueList: + issues = self.manager.change(action, description) + issues.hint = "Use the .inspect.issues() for more details." + # Todo Depending on the action, change the issue.action to reflect the action taken. + return issues + def rule_transform(self, *transformer: VerifiedRulesTransformer) -> IssueList: if not transformer: raise NeatSessionError("No transformers provided.") diff --git a/cognite/neat/_store/__init__.py b/cognite/neat/_store/__init__.py index ac08a72fc..060b66de3 100644 --- a/cognite/neat/_store/__init__.py +++ b/cognite/neat/_store/__init__.py @@ -1,4 +1,5 @@ from ._graph_store import NeatGraphStore +from ._manager import NeatStoreManager from ._rules_store import NeatRulesStore -__all__ = ["NeatGraphStore", "NeatRulesStore"] +__all__ = ["NeatGraphStore", "NeatRulesStore", "NeatStoreManager"] diff --git a/cognite/neat/_store/_manager.py b/cognite/neat/_store/_manager.py new file mode 100644 index 000000000..f06e42d60 --- /dev/null +++ b/cognite/neat/_store/_manager.py @@ -0,0 +1,45 @@ +from typing import TypeAlias + +from cognite.neat._graph.extractors import BaseExtractor +from cognite.neat._graph.transformers import BaseTransformer, BaseTransformerStandardised +from cognite.neat._issues import IssueList +from cognite.neat._rules.importers import BaseImporter +from cognite.neat._rules.transformers import VerifiedRulesTransformer + +from ._graph_store import NeatGraphStore +from ._rules_store import NeatRulesStore +from .exceptions import InvalidAction + +Action: TypeAlias = ( + BaseExtractor | BaseTransformerStandardised | BaseTransformer | BaseImporter | VerifiedRulesTransformer +) + + +class NeatStoreManager: + def __init__(self, instances: NeatGraphStore, rules: NeatRulesStore) -> None: + self._instances = instances + self._rules = rules + + def change(self, action: Action, description: str | None = None) -> IssueList: + """Perform an action on the state of either the rule or instance store.""" + if error_message := self._can_perform(action): + raise InvalidAction(description or action.description, error_message) # type: ignore[union-attr] + return self._perform(action) + + def _can_perform(self, action: Action) -> str: + if isinstance(action, BaseExtractor) and not self._rules.empty: + return "Cannot extract instances when a data model is in the session. You need to restart the session." + raise NotImplementedError() + + def _perform(self, action: Action) -> IssueList: + match action: + case _ if isinstance(action, BaseExtractor): + return self._instances.write(action) + case _ if isinstance(action, BaseTransformerStandardised | BaseTransformer): + return self._instances.transform(action) + case _ if isinstance(action, BaseImporter): + return self._rules.import_rules(action) + case _ if isinstance(action, VerifiedRulesTransformer): + return self._rules.transform(action) + case _: + raise NotImplementedError() diff --git a/cognite/neat/_store/exceptions.py b/cognite/neat/_store/exceptions.py index 2ef0fdd36..947ff9188 100644 --- a/cognite/neat/_store/exceptions.py +++ b/cognite/neat/_store/exceptions.py @@ -57,3 +57,12 @@ class EmptyStore(NeatStoreError, RuntimeError): """Raised when the store is empty""" ... + + +class InvalidAction(NeatStoreError, RuntimeError): + def __init__(self, action: str, error_message: str) -> None: + self.action = action + self.error_message = error_message + + def __str__(self) -> str: + return f"Cannot do {self.action}: {self.error_message}"