[NEAT-478] 🚅 Load Engine (#722)

* refactor: shell for load neat engine

* build; added packaging dependency

* feat: wrote loading

* fix: robustify

* build: extend test runner

* style: better name

* build; fix dependencies

* build: fix image

Makefile

@@ -39,7 +39,7 @@ start-ui-dev:
 
 poetry-export:
 	@echo "Exporting poetry dependencies"
-	poetry export -f requirements.txt --output requirements.txt --extras "graphql" --extras "service"
+	poetry export -f requirements.txt --output requirements.txt --extras "service"
 
 build-docker: poetry-export
 	@echo "Building docker image"

cognite/neat/_session/_base.py

@@ -20,6 +20,7 @@
     Change,
 )
 
+from ._engine import load_neat_engine
 from ._inspect import InspectAPI
 from ._prepare import PrepareAPI
 from ._read import ReadAPI
@@ -37,6 +38,7 @@ def __init__(
         client: CogniteClient | None = None,
         storage: Literal["memory", "oxigraph"] = "memory",
         verbose: bool = True,
+        load_engine: Literal["newest", "cache", "skip"] = "cache",
     ) -> None:
         self._client = client
         self._verbose = verbose
@@ -47,6 +49,8 @@ def __init__(
         self.show = ShowAPI(self._state)
         self.set = SetAPI(self._state, verbose)
         self.inspect = InspectAPI(self._state)
+        if load_engine != "skip" and (engine_version := load_neat_engine(client, load_engine)):
+            print(f"Neat Engine {engine_version} loaded.")
 
     @property
     def version(self) -> str:

cognite/neat/_session/_engine.py

@@ -0,0 +1,120 @@
+import os
+import re
+import shutil
+import sys
+import tempfile
+import warnings
+from collections.abc import Callable
+from pathlib import Path
+from typing import Literal, cast
+
+from cognite.client import CogniteClient
+from packaging.version import Version
+from packaging.version import parse as parse_version
+
+from cognite.neat._issues.errors import NeatValueError
+from cognite.neat._version import __engine__
+
+ENVIRONMENT_VARIABLE = "NEATENGINE"
+PACKAGE_NAME = "neatengine"
+PYVERSION = f"{sys.version_info.major}{sys.version_info.minor}"
+
+
+def load_neat_engine(client: CogniteClient | None, location: Literal["newest", "cache"]) -> str | None:
+    if location not in ["newest", "cache"]:
+        raise NeatValueError(f"Cannot load engine from location: {location}")
+    cache_dir = Path(tempfile.gettempdir()) / PACKAGE_NAME
+    cache_dir.mkdir(exist_ok=True)
+    pattern = re.compile(rf"{PACKAGE_NAME}-(\d+\.\d+\.\d+)-{PYVERSION}.zip")
+
+    candidates: dict[Version, Callable[[], Path]] = {}
+    if location == "cache" and cache_dir.exists():
+        candidates = _load_from_path(cache_dir, pattern)
+
+    if location == "newest" or not candidates:
+        # Loading in revrse order of priority
+        # 3. Downloads folder
+        candidates = _load_from_path(Path.home() / "Downloads", pattern)
+        # 2. CDF
+        if client:
+            candidates.update(_load_from_cdf(client, pattern, cache_dir))
+        # 1. Environment variable
+        if ENVIRONMENT_VARIABLE in os.environ:
+            environ_path = Path(os.environ[ENVIRONMENT_VARIABLE])
+            if environ_path.exists():
+                candidates.update(_load_from_path(environ_path, pattern))
+            else:
+                warnings.warn(
+                    f"Environment variable {ENVIRONMENT_VARIABLE} points to non-existing path: {environ_path}",
+                    UserWarning,
+                    stacklevel=2,
+                )
+
+    if not candidates:
+        return None
+
+    if not __engine__.startswith("^"):
+        # Using value error as this is a developer error
+        raise ValueError(f"Invalid engine version: {__engine__}")
+
+    lower_bound = parse_version(__engine__[1:])
+    upper_bound = Version(f"{lower_bound.major+1}.0.0")
+    selected_version = max(
+        (version for version in candidates.keys() if lower_bound <= version < upper_bound), default=None
+    )
+    if not selected_version:
+        return None
+    source_path = candidates[selected_version]()
+    destination_path = cache_dir / source_path.name
+    if not destination_path.exists():
+        shutil.copy(source_path, destination_path)
+    sys.path.append(str(destination_path))
+    try:
+        from neatengine._version import __version__ as engine_version  # type: ignore[import-not-found]
+    except ImportError:
+        return None
+    return engine_version
+
+
+def _load_from_path(path: Path, pattern) -> dict[Version, Callable[[], Path]]:
+    if path.is_file() and (match := pattern.match(path.name)):
+        return {parse_version(match.group(1)): lambda: path}
+    elif path.is_dir():
+        output: dict[Version, Callable[[], Path]] = {}
+        for candidate in path.iterdir():
+            if candidate.is_file() and (match := pattern.match(candidate.name)):
+                # Setting default value to ensure we use the candidate from the current iteration
+                # If not set, the function will use the last candidate from the loop
+                def return_path(the_path: Path = candidate) -> Path:
+                    return the_path
+
+                output[parse_version(match.group(1))] = return_path
+
+        return output
+    return {}
+
+
+def _load_from_cdf(
+    client: CogniteClient, pattern: re.Pattern[str], cache_dir: Path
+) -> dict[Version, Callable[[], Path]]:
+    file_metadata = client.files.list(
+        limit=-1,
+        data_set_external_ids=PACKAGE_NAME,
+        external_id_prefix=PACKAGE_NAME,
+        metadata={"python_version": PYVERSION},
+    )
+    output: dict[Version, Callable[[], Path]] = {}
+    for file in file_metadata:
+        name = cast(str, file.name)
+
+        # Use function to lazily download file
+        # Setting default value to ensure we use the file_id from the current iteration
+        # If not set, the function will use the last file_id from the loop
+        def download_file(file_id: int = file.id, filename: str = name) -> Path:
+            client.files.download(cache_dir, file_id)
+            return cache_dir / filename
+
+        if match := pattern.match(name):
+            output[parse_version(match.group(1))] = download_file
+
+    return output

cognite/neat/_version.py

@@ -1 +1,2 @@
 __version__ = "0.96.6"
+__engine__ = "^0.1.0"