cognite-neat 0.125.1__py3-none-any.whl → 0.126.1__py3-none-any.whl

cognite/neat/_store/_store.py

@@ -0,0 +1,124 @@
+from collections import UserList
+from collections.abc import Callable
+from datetime import datetime, timezone
+from typing import Any, cast
+
+from cognite.neat._data_model._shared import OnSuccess, OnSuccessIssuesChecker, OnSuccessResultProducer
+from cognite.neat._data_model.exporters 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._exceptions import DataModelImportException
+from cognite.neat._issues import IssueList
+from cognite.neat._state_machine._states import EmptyState, State
+
+from ._provenance import Change, Provenance
+
+Agents = DMSTableExporter | DMSTableImporter | DMSImporter
+
+
+class NeatStore:
+    def __init__(self) -> None:
+        self.physical_data_model = DataModelList()
+        self.provenance = Provenance()
+        self.state: State = EmptyState()
+
+    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
+            # in case of read of data model result will be same as target entity
+            change.result = change.target_entity
+
+        self.provenance.append(change)
+
+    def write_physical(self, writer: DMSTableExporter, on_success: OnSuccess | None = None, **kwargs: Any) -> None:
+        """Write object into the store"""
+        self._can_agent_do_activity(writer)
+
+        change, _ = self._do_activity(writer.export, on_success, data_model=self.physical_data_model[-1], **kwargs)
+
+        if not change.issues:
+            change.target_entity = "ExternalEntity"
+            self.state = self.state.transition(writer)
+            change.target_state = self.state
+
+        self.provenance.append(change)
+
+    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):
+            raise RuntimeError(f"Cannot run {type(agent).__name__} in state {self.state}")
+
+        # 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)
+        result: PhysicalDataModel | None = None
+        issues = IssueList()
+        errors = IssueList()
+
+        try:
+            result = activity(**kwargs)
+
+            if result and on_success:
+                if isinstance(on_success, OnSuccessIssuesChecker):
+                    on_success.run(result)
+                    issues.extend(on_success.issues)
+                elif isinstance(on_success, OnSuccessResultProducer):
+                    raise NotImplementedError("OnSuccessResultProducer is not implemented yet.")
+                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,
+            activity=Change.standardize_activity_name(activity.__name__, start, end),
+        ), result
+
+
+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/http_client/__init__.py

@@ -1,45 +1,39 @@
 from ._client import HTTPClient
 from ._data_classes import (
-    FailedItem,
-    FailedRequestItem,
+    ErrorDetails,
+    FailedRequestItems,
     FailedRequestMessage,
     FailedResponse,
+    FailedResponseItems,
     HTTPMessage,
-    ItemIDMessage,
+    ItemBody,
+    ItemIDBody,
     ItemMessage,
-    ItemResponse,
     ItemsRequest,
-    MissingItem,
-    ParamRequest,
+    ParametersRequest,
     RequestMessage,
     ResponseMessage,
     SimpleBodyRequest,
-    SuccessItem,
     SuccessResponse,
-    UnexpectedItem,
-    UnknownRequestItem,
-    UnknownResponseItem,
+    SuccessResponseItems,
 )
 
 __all__ = [
-    "FailedItem",
-    "FailedRequestItem",
+    "ErrorDetails",
+    "FailedRequestItems",
     "FailedRequestMessage",
     "FailedResponse",
+    "FailedResponseItems",
     "HTTPClient",
     "HTTPMessage",
-    "ItemIDMessage",
+    "ItemBody",
+    "ItemIDBody",
     "ItemMessage",
-    "ItemResponse",
     "ItemsRequest",
-    "MissingItem",
-    "ParamRequest",
+    "ParametersRequest",
     "RequestMessage",
     "ResponseMessage",
     "SimpleBodyRequest",
-    "SuccessItem",
     "SuccessResponse",
-    "UnexpectedItem",
-    "UnknownRequestItem",
-    "UnknownResponseItem",
+    "SuccessResponseItems",
 ]

cognite/neat/_utils/http_client/_client.py

@@ -8,20 +8,21 @@ from typing import Literal
 
 import httpx
 from cognite.client import ClientConfig, global_config
-from cognite.client.utils import _json
 
 from cognite.neat._utils.auxiliary import get_current_neat_version
 from cognite.neat._utils.http_client._config import get_user_agent
 from cognite.neat._utils.http_client._data_classes import (
+    APIResponse,
     BodyRequest,
+    ErrorDetails,
     FailedRequestMessage,
     HTTPMessage,
     ItemsRequest,
-    ParamRequest,
+    ParametersRequest,
     RequestMessage,
     ResponseMessage,
 )
-from cognite.neat._utils.useful_types import JsonVal
+from cognite.neat._utils.useful_types import PrimaryTypes
 
 if sys.version_info >= (3, 11):
     from typing import Self
@@ -98,7 +99,7 @@ class HTTPClient:
             results = self._handle_error(e, message)
         return results
 
-    def request_with_retries(self, message: RequestMessage) -> Sequence[ResponseMessage | FailedRequestMessage]:
+    def request_with_retries(self, message: RequestMessage) -> APIResponse:
         """Send an HTTP request and handle retries.
 
         This method will keep retrying the request until it either succeeds or
@@ -118,7 +119,7 @@ class HTTPClient:
             raise RuntimeError(f"RequestMessage has already been attempted {message.total_attempts} times.")
         pending_requests: deque[RequestMessage] = deque()
         pending_requests.append(message)
-        final_responses: list[ResponseMessage | FailedRequestMessage] = []
+        final_responses = APIResponse()
 
         while pending_requests:
             current_request = pending_requests.popleft()
@@ -157,36 +158,16 @@ class HTTPClient:
             headers["Content-Encoding"] = "gzip"
         return headers
 
-    @staticmethod
-    def _prepare_payload(item: BodyRequest) -> str | bytes:
-        """
-        Prepare the payload for the HTTP request.
-        This method should be overridden in subclasses to customize the payload format.
-        """
-        data: str | bytes
-        try:
-            data = _json.dumps(item.body(), allow_nan=False)
-        except ValueError as e:
-            # A lot of work to give a more human friendly error message when nans and infs are present:
-            msg = "Out of range float values are not JSON compliant"
-            if msg in str(e):  # exc. might e.g. contain an extra ": nan", depending on build (_json.make_encoder)
-                raise ValueError(f"{msg}. Make sure your data does not contain NaN(s) or +/- Inf!").with_traceback(
-                    e.__traceback__
-                ) from None
-            raise
-
-        if not global_config.disable_gzip:
-            data = gzip.compress(data.encode())
-        return data
-
     def _make_request(self, item: RequestMessage) -> httpx.Response:
         headers = self._create_headers(item.api_version)
-        params: dict[str, str] | None = None
-        if isinstance(item, ParamRequest):
+        params: dict[str, PrimaryTypes] | None = None
+        if isinstance(item, ParametersRequest):
             params = item.parameters
         data: str | bytes | None = None
         if isinstance(item, BodyRequest):
-            data = self._prepare_payload(item)
+            data = item.data()
+            if not global_config.disable_gzip:
+                data = gzip.compress(data.encode("utf-8"))
         return self.session.request(
             method=item.method,
             url=item.endpoint_url,
@@ -202,21 +183,12 @@ class HTTPClient:
         response: httpx.Response,
         request: RequestMessage,
     ) -> Sequence[HTTPMessage]:
-        try:
-            body = response.json()
-        except ValueError as e:
-            return request.create_responses(response, error_message=f"Invalid JSON response: {e!s}")
-
-        error_obj = body.get("error", {})
-        is_auto_retryable = False
-        if isinstance(error_obj, dict):
-            is_auto_retryable = error_obj.get("isAutoRetryable", False)
-
         if 200 <= response.status_code < 300:
-            return request.create_responses(response, body)
-        elif (
+            return request.create_success_response(response)
+
+        if (
             isinstance(request, ItemsRequest)
-            and len(request.items) > 1
+            and len(request.body.items) > 1
             and response.status_code in self._split_items_status_codes
         ):
             # 4XX: Status there is at least one item that is invalid, split the batch to get all valid items processed
@@ -226,31 +198,20 @@ class HTTPClient:
                 status_attempts += 1
             splits = request.split(status_attempts=status_attempts)
             if splits[0].tracker and splits[0].tracker.limit_reached():
-                return request.create_responses(response, body, self._get_error_message(body, response.text))
+                return request.create_failure_response(response)
             return splits
-        elif request.status_attempt < self._max_retries and (
-            response.status_code in self._retry_status_codes or is_auto_retryable
+
+        error = ErrorDetails.from_response(response)
+
+        if request.status_attempt < self._max_retries and (
+            response.status_code in self._retry_status_codes or error.is_auto_retryable
         ):
             request.status_attempt += 1
             time.sleep(self._backoff_time(request.total_attempts))
             return [request]
         else:
             # Permanent failure
-            return request.create_responses(response, body, self._get_error_message(body, response.text))
-
-    @staticmethod
-    def _get_error_message(body: JsonVal, default: str) -> str:
-        error = default
-        if not isinstance(body, dict):
-            return error
-        if "error" not in body:
-            return error
-        error_nested = body["error"]
-        if isinstance(error_nested, str):
-            return error_nested
-        if isinstance(error_nested, dict) and "message" in error_nested and isinstance(error_nested["message"], str):
-            return error_nested["message"]
-        return error
+            return request.create_failure_response(response)
 
     @staticmethod
     def _backoff_time(attempts: int) -> float: