cognite-neat 0.123.26__py3-none-any.whl → 1.0.22__py3-none-any.whl

cognite/neat/_store/_store.py

@@ -0,0 +1,220 @@
+from collections import UserList
+from collections.abc import Callable
+from datetime import datetime, timezone
+from typing import Any, cast
+
+from cognite.neat._client.client import NeatClient
+from cognite.neat._config import NeatConfig
+from cognite.neat._data_model._shared import OnSuccess, OnSuccessIssuesChecker, OnSuccessResultProducer
+from cognite.neat._data_model._snapshot import SchemaSnapshot
+from cognite.neat._data_model.deployer.data_classes import DeploymentResult
+from cognite.neat._data_model.deployer.deployer import SchemaDeployer
+from cognite.neat._data_model.exporters import DMSExporter, DMSFileExporter
+from cognite.neat._data_model.exporters._api_exporter import DMSAPIExporter
+from cognite.neat._data_model.exporters._table_exporter.exporter import DMSTableExporter
+from cognite.neat._data_model.importers import DMSImporter, DMSTableImporter
+from cognite.neat._data_model.models.dms import RequestSchema as PhysicalDataModel
+from cognite.neat._data_model.models.dms._limits import SchemaLimits
+from cognite.neat._exceptions import DataModelImportException
+from cognite.neat._issues import IssueList
+from cognite.neat._state_machine._states import EmptyState, PhysicalState, State
+from cognite.neat._utils.text import NEWLINE
+
+from ._provenance import Change, Provenance
+
+Agents = DMSExporter | DMSTableImporter | DMSImporter
+
+
+class NeatStore:
+    def __init__(self, config: NeatConfig, client: NeatClient) -> None:
+        self.physical_data_model = DataModelList()
+        self.provenance = Provenance()
+        self.state: State = EmptyState()
+        self._client = client
+        self._config = config
+
+        # Placeholder for CDF schema and limit snapshot
+        self._cdf_snapshot: SchemaSnapshot | None = None
+        self._cdf_limits: SchemaLimits | None = None
+
+    @property
+    def cdf_limits(self) -> SchemaLimits:
+        if not self._cdf_limits:
+            self._cdf_limits = SchemaLimits.from_api_response(self._client.statistics.project())
+        return self._cdf_limits
+
+    @property
+    def cdf_snapshot(self) -> SchemaSnapshot:
+        if not self._cdf_snapshot:
+            self._cdf_snapshot = SchemaSnapshot.fetch_entire_cdf(self._client)
+        return self._cdf_snapshot
+
+    def read_physical(self, reader: DMSImporter, on_success: OnSuccess | None = None) -> None:
+        """Read object from the store"""
+        self._can_agent_do_activity(reader)
+
+        change, data_model = self._do_activity(reader.to_data_model, on_success)
+
+        if data_model:
+            change.target_entity = self.physical_data_model.generate_reference(cast(PhysicalDataModel, data_model))
+            self.physical_data_model.append(data_model)
+            self.state = self.state.transition(reader)
+            change.target_state = self.state
+
+        self.provenance.append(change)
+
+    def write_physical(self, writer: DMSExporter, on_success: OnSuccess | None = None, **kwargs: Any) -> None:
+        """Write object into the store"""
+        self._can_agent_do_activity(writer)
+
+        activity: Callable
+        if isinstance(writer, DMSFileExporter):
+            activity = writer.export_to_file
+            if not kwargs.get("file_path"):
+                raise RuntimeError("file_path must be provided when using a DMSFileExporter")
+        else:
+            activity = writer.export
+
+        data_model = self._gather_data_model(writer)
+
+        change, _ = self._do_activity(activity, on_success, data_model=data_model, **kwargs)
+
+        if not change.issues:
+            change.target_entity = "ExternalEntity"
+            self.state = self.state.transition(writer)
+            change.target_state = self.state
+
+        self.provenance.append(change)
+
+        if (
+            isinstance(writer, DMSAPIExporter)
+            and isinstance(on_success, SchemaDeployer)
+            and not on_success.options.dry_run
+        ):
+            # Update CDF snapshot after successful deployment
+            self._cdf_snapshot = SchemaSnapshot.fetch_entire_cdf(self._client)
+
+    def _gather_data_model(self, writer: DMSExporter) -> PhysicalDataModel:
+        """Gather the current physical data model from the store
+
+        Args:
+            writer (DMSExporter): The exporter that will be used to write the data model.
+        """
+        # getting provenance of the last successful physical data model read
+        change = self.provenance.last_physical_data_model_read()
+
+        if not change:
+            raise RuntimeError("No successful physical data model read found in provenance.")
+
+        # We do not want to modify the data model for API representations
+        if not isinstance(writer, DMSTableExporter):
+            return self.physical_data_model[-1]
+
+        # This will handle data model that are partially and require to be converted to
+        # tabular representation to include all containers referenced by views.
+        copy = self.physical_data_model[-1].model_copy(deep=True)
+        container_refs = {container.as_reference() for container in copy.containers}
+
+        for view in copy.views:
+            for container in view.used_containers:
+                if container not in container_refs and (cdf_container := self.cdf_snapshot.containers.get(container)):
+                    copy.containers.append(cdf_container)
+                    container_refs.add(container)
+
+        return copy
+
+    def _can_agent_do_activity(self, agent: Agents) -> None:
+        """Validate if activity can be performed in the current state and considering provenance"""
+        if not self.state.can_transition(agent):
+            # specific error messages for common mistakes
+            if isinstance(agent, DMSImporter) and isinstance(self.state, PhysicalState):
+                raise RuntimeError(
+                    "⚠️ Cannot read data model, there is already a data model in the session!"
+                    f"{NEWLINE}Start a new session to read a new data model."
+                )
+
+            if isinstance(agent, DMSExporter) and isinstance(self.state, EmptyState):
+                raise RuntimeError(
+                    "⚠️ Cannot write data model, there is no data model in the session!"
+                    f"{NEWLINE}Read a data model first!"
+                )
+            raise RuntimeError(f"Cannot run {type(agent).__name__} in state {self.state}")
+
+        if (
+            isinstance(agent, DMSExporter)
+            and self.provenance.last_change
+            and (error_count := self.provenance.last_change.error_count) > 0
+        ):
+            raise RuntimeError(
+                f"⚠️ Cannot write data model, the model has {error_count} errors!"
+                f"{NEWLINE}Resolve issues before exporting the data model."
+                f"{NEWLINE}You can inspect issues using neat.issues"
+            )
+
+        # need implementation of checking if required predecessor activities have been done
+        # this will be done by running self.provenance.can_agent_do_activity(agent)
+
+    def _do_activity(
+        self, activity: Callable, on_success: OnSuccess | None = None, **kwargs: Any
+    ) -> tuple[Change, PhysicalDataModel | None]:
+        """Execute activity and capture timing, results, and issues"""
+        start = datetime.now(timezone.utc)
+        created_data_model: PhysicalDataModel | None = None
+        issues = IssueList()
+        errors = IssueList()
+        deployment_result: DeploymentResult | None = None
+
+        try:
+            created_data_model = activity(**kwargs)
+            if created_data_model and on_success:
+                on_success.run(created_data_model)
+                if isinstance(on_success, OnSuccessIssuesChecker):
+                    issues.extend(on_success.issues)
+                elif isinstance(on_success, OnSuccessResultProducer):
+                    deployment_result = on_success.result
+                else:
+                    raise RuntimeError(f"Unknown OnSuccess type {type(on_success).__name__}")
+
+        # we catch import exceptions to capture issues and errors in provenance
+        except DataModelImportException as e:
+            errors.extend(e.errors)
+
+        # these are all other errors, such as missing file, wrong format, etc.
+        except Exception as e:
+            raise e
+
+        end = datetime.now(timezone.utc)
+
+        return Change(
+            start=start,
+            end=end,
+            source_state=self.state,
+            agent=type(activity.__self__).__name__ if hasattr(activity, "__self__") else "UnknownAgent",
+            issues=issues,
+            errors=errors,
+            result=deployment_result,
+            activity=Change.standardize_activity_name(activity.__name__, start, end),
+        ), created_data_model
+
+
+class DataModelList(UserList[PhysicalDataModel]):
+    def iteration(self, data_model: PhysicalDataModel) -> int:
+        """Get iteration number for data model"""
+        for i, existing in enumerate(self):
+            if existing.data_model == data_model.data_model:
+                return i + 2
+        return 1
+
+    def generate_reference(self, data_model: PhysicalDataModel) -> str:
+        """Generate reference string for data model based on iteration"""
+        space = data_model.data_model.space
+        external_id = data_model.data_model.external_id
+        version = data_model.data_model.version
+        iteration = self.iteration(data_model)
+
+        return f"physical/{space}/{external_id}/{version}/{iteration}"
+
+    def get_by_reference(self, reference: str) -> PhysicalDataModel | None:
+        """Get data model by reference string"""
+
+        raise NotImplementedError("Not implemented yet")

cognite/neat/_utils/_reader.py

@@ -0,0 +1,194 @@
+import tempfile
+from abc import ABC, abstractmethod
+from collections.abc import Iterable
+from io import StringIO
+from pathlib import Path
+from typing import IO, Any, TextIO
+from urllib.parse import urlparse
+
+import requests
+
+
+class NeatReader(ABC):
+    @classmethod
+    def create(cls, io: Any) -> "NeatReader":
+        if isinstance(io, str):
+            url = urlparse(io)
+            if url.scheme == "https" and url.netloc.endswith("github.com"):
+                return GitHubReader(io)
+            elif url.scheme == "https":
+                return HttpFileReader(io, url.path)
+
+        if isinstance(io, str):
+            return PathReader(Path(io))
+        if isinstance(io, Path):
+            return PathReader(io)
+        raise ValueError(f"Unsupported type: {type(io)}")
+
+    @property
+    def name(self) -> str:
+        return str(self)
+
+    @abstractmethod
+    def read_text(self) -> str:
+        """Read the buffer as a string"""
+        raise NotImplementedError()
+
+    @abstractmethod
+    def read_bytes(self) -> bytes:
+        """Read the buffer as bytes"""
+        raise NotImplementedError()
+
+    @abstractmethod
+    def size(self) -> int:
+        """Size of the buffer in bytes"""
+        raise NotImplementedError()
+
+    @abstractmethod
+    def iterate(self, chunk_size: int) -> Iterable[str]:
+        """Iterate over the buffer in chunks
+
+        Args:
+            chunk_size: Size of each chunk in bytes
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def __enter__(self) -> IO:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def __str__(self) -> str:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def exists(self) -> bool:
+        raise NotImplementedError
+
+    @abstractmethod
+    def materialize_path(self) -> Path:
+        raise NotImplementedError
+
+
+class PathReader(NeatReader):
+    def __init__(self, path: Path):
+        self.path = path
+        self._io: TextIO | None = None
+
+    @property
+    def name(self) -> str:
+        return self.path.name
+
+    def read_text(self) -> str:
+        return self.path.read_text()
+
+    def read_bytes(self) -> bytes:
+        return self.path.read_bytes()
+
+    def size(self) -> int:
+        return self.path.stat().st_size
+
+    def iterate(self, chunk_size: int) -> Iterable[str]:
+        with self.path.open(mode="r") as f:
+            while chunk := f.read(chunk_size):
+                yield chunk
+
+    def __enter__(self) -> TextIO:
+        file = self.path.open(mode="r")
+        self._io = file
+        return file
+
+    def __exit__(self) -> None:
+        if self._io:
+            self._io.close()
+
+    def __str__(self) -> str:
+        return self.path.as_posix()
+
+    def exists(self) -> bool:
+        return self.path.exists()
+
+    def materialize_path(self) -> Path:
+        return self.path
+
+
+class HttpFileReader(NeatReader):
+    def __init__(self, url: str, path: str):
+        self._url = url
+        self.path = path
+
+    @property
+    def name(self) -> str:
+        if "/" in self.path:
+            return self.path.rsplit("/", maxsplit=1)[-1]
+        return self.path
+
+    def read_text(self) -> str:
+        response = requests.get(self._url)
+        response.raise_for_status()
+        return response.text
+
+    def read_bytes(self) -> bytes:
+        response = requests.get(self._url)
+        response.raise_for_status()
+        return response.content
+
+    def size(self) -> int:
+        response = requests.head(self._url)
+        response.raise_for_status()
+        return int(response.headers["Content-Length"])
+
+    def iterate(self, chunk_size: int) -> Iterable[str]:
+        with requests.get(self._url, stream=True) as response:
+            response.raise_for_status()
+            for chunk in response.iter_content(chunk_size):
+                yield chunk.decode("utf-8")
+
+    def __enter__(self) -> IO:
+        return StringIO(self.read_text())
+
+    def __str__(self) -> str:
+        return self._url
+
+    def exists(self) -> bool:
+        response = requests.head(self._url)
+        return 200 <= response.status_code < 400
+
+    def materialize_path(self) -> Path:
+        path = Path(tempfile.gettempdir()).resolve() / "neat" / self.name
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_bytes(self.read_bytes())
+        return path
+
+
+class GitHubReader(HttpFileReader):
+    raw_url = "https://raw.githubusercontent.com/"
+
+    def __init__(self, raw: str):
+        self.raw = raw
+        self.repo, path = self._parse_url(raw)
+        super().__init__(f"{self.raw_url}{self.repo}/main/{path}", path)
+
+    @staticmethod
+    def _parse_url(url: str) -> tuple[str, str]:
+        parsed = urlparse(url)
+        if parsed.scheme != "https":
+            raise ValueError(f"Unsupported scheme: {parsed.scheme}")
+
+        path = parsed.path.lstrip("/")
+        if parsed.netloc == "github.com":
+            repo, path = path.split("/blob/main/", maxsplit=1)
+            return repo, path
+
+        elif parsed.netloc == "api.github.com":
+            repo, path = path.removeprefix("repos/").split("/contents/", maxsplit=1)
+            return repo, path
+
+        elif parsed.netloc == "raw.githubusercontent.com":
+            repo, path = path.split("/main/", maxsplit=1)
+            return repo, path
+
+        raise ValueError(f"Unsupported netloc: {parsed.netloc}")
+
+    def __str__(self) -> str:
+        return self.raw

cognite/neat/_utils/auxiliary.py

@@ -0,0 +1,49 @@
+import importlib
+import inspect
+from abc import ABC
+from types import ModuleType
+from typing import TypeVar
+
+from cognite.neat import _version
+from cognite.neat._exceptions import NeatImportError
+
+
+def get_current_neat_version() -> str:
+    return _version.__version__
+
+
+T_Cls = TypeVar("T_Cls")
+
+
+def get_concrete_subclasses(base_cls: type[T_Cls], exclude_direct_abc_inheritance: bool = True) -> list[type[T_Cls]]:
+    """
+    Returns a list of all concrete subclasses of the given base class.
+    Args:
+        base_cls (type[T_Cls]): The base class to find subclasses for.
+        exclude_direct_abc_inheritance (bool): If True, excludes classes that directly inherit from `abc.ABC`.
+            This is used as a marker to filter out intermediate base classes. Defaults to True.
+    Returns:
+        list[type[T_Cls]]: A list of concrete subclasses of the base class.
+    """
+    to_check = [base_cls]
+    subclasses: list[type[T_Cls]] = []
+    seen: set[type[T_Cls]] = {base_cls}
+    while to_check:
+        current_cls = to_check.pop()
+        for subclass in current_cls.__subclasses__():
+            if subclass in seen:
+                continue
+            if (not inspect.isabstract(subclass)) and (
+                not exclude_direct_abc_inheritance or ABC not in subclass.__bases__
+            ):
+                subclasses.append(subclass)
+            seen.add(subclass)
+            to_check.append(subclass)
+    return subclasses
+
+
+def local_import(module: str, extra: str) -> ModuleType:
+    try:
+        return importlib.import_module(module)
+    except ImportError as e:
+        raise NeatImportError(module.split(".")[0], extra) from e

cognite/neat/_utils/collection.py

@@ -0,0 +1,11 @@
+from collections.abc import Iterator, Sequence
+from typing import TypeVar
+
+T_Sequence = TypeVar("T_Sequence", bound=Sequence)
+
+
+def chunker_sequence(sequence: T_Sequence, size: int) -> Iterator[T_Sequence]:
+    """Yield successive n-sized chunks from sequence."""
+    for i in range(0, len(sequence), size):
+        # MyPy does not expect sequence[i : i + size] to be of type T_Sequence
+        yield sequence[i : i + size]  # type: ignore[misc]

cognite/neat/_utils/http_client/__init__.py

@@ -0,0 +1,39 @@
+from ._client import HTTPClient
+from ._data_classes import (
+    ErrorDetails,
+    FailedRequestItems,
+    FailedRequestMessage,
+    FailedResponse,
+    FailedResponseItems,
+    HTTPMessage,
+    ItemBody,
+    ItemIDBody,
+    ItemMessage,
+    ItemsRequest,
+    ParametersRequest,
+    RequestMessage,
+    ResponseMessage,
+    SimpleBodyRequest,
+    SuccessResponse,
+    SuccessResponseItems,
+)
+
+__all__ = [
+    "ErrorDetails",
+    "FailedRequestItems",
+    "FailedRequestMessage",
+    "FailedResponse",
+    "FailedResponseItems",
+    "HTTPClient",
+    "HTTPMessage",
+    "ItemBody",
+    "ItemIDBody",
+    "ItemMessage",
+    "ItemsRequest",
+    "ParametersRequest",
+    "RequestMessage",
+    "ResponseMessage",
+    "SimpleBodyRequest",
+    "SuccessResponse",
+    "SuccessResponseItems",
+]