cognite-neat 1.0.24__tar.gz → 1.0.26__tar.gz

{cognite_neat-1.0.24 → cognite_neat-1.0.26}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognite-neat
-Version: 1.0.24
+Version: 1.0.26
 Summary: Knowledge graph transformation
 Author: Nikola Vasiljevic, Anders Albert
 Author-email: Nikola Vasiljevic <nikola.vasiljevic@cognite.com>, Anders Albert <anders.albert@cognite.com>

{cognite_neat-1.0.24 → cognite_neat-1.0.26}/cognite/neat/__init__.py

@@ -1,4 +1,4 @@
-from cognite.neat._v0.core._utils.auth import get_cognite_client
+from cognite.neat._client.init.main import get_cognite_client
 
 from ._config import NeatConfig, get_neat_config_from_file
 from ._session import NeatSession

cognite_neat-1.0.26/cognite/neat/_client/api.py

@@ -0,0 +1,178 @@
+from abc import ABC, abstractmethod
+from collections.abc import Iterable, Sequence
+from dataclasses import dataclass
+from typing import Any, Generic, Literal, TypeAlias, TypeVar
+
+from pydantic import BaseModel, JsonValue, TypeAdapter
+
+from cognite.neat._client.config import NeatClientConfig
+from cognite.neat._data_model.models.dms._base import T_Resource, T_Response
+from cognite.neat._utils.collection import chunker_sequence
+from cognite.neat._utils.http_client import HTTPClient, ParametersRequest, SimpleBodyRequest, SuccessResponse
+from cognite.neat._utils.useful_types import T_Reference
+
+from .data_classes import PagedResponse
+
+_T_BaseModel = TypeVar("_T_BaseModel", bound=BaseModel)
+
+
+@dataclass(frozen=True)
+class Endpoint:
+    method: Literal["GET", "POST"]
+    path: str
+    item_limit: int
+
+
+APIMethod: TypeAlias = Literal["apply", "retrieve", "delete", "list"]
+
+
+class NeatAPI(Generic[T_Reference, T_Resource, T_Response], ABC):
+    def __init__(
+        self, neat_config: NeatClientConfig, http_client: HTTPClient, endpoint_map: dict[APIMethod, Endpoint]
+    ) -> None:
+        self._config = neat_config
+        self._http_client = http_client
+        self._method_endpoint_map = endpoint_map
+
+    @abstractmethod
+    def _validate_page_response(self, response: SuccessResponse) -> PagedResponse[T_Response]:
+        """Parse a single item response."""
+        raise NotImplementedError()
+
+    @abstractmethod
+    def _validate_id_response(self, response: SuccessResponse) -> list[T_Reference]:
+        """Parse a single item response."""
+        raise NotImplementedError()
+
+    def _make_url(self, path: str = "") -> str:
+        """Create the full URL for this resource endpoint."""
+        return self._config.create_api_url(path)
+
+    def _request_item_response(
+        self,
+        items: Sequence[BaseModel],
+        method: APIMethod,
+        extra_body: dict[str, Any] | None = None,
+    ) -> list[T_Response]:
+        response_items: list[T_Response] = []
+        for response in self._chunk_requests(items, method, extra_body):
+            response_items.extend(self._validate_page_response(response).items)
+        return response_items
+
+    def _request_id_response(
+        self,
+        items: Sequence[BaseModel],
+        method: APIMethod,
+        extra_body: dict[str, Any] | None = None,
+    ) -> list[T_Reference]:
+        response_items: list[T_Reference] = []
+        for response in self._chunk_requests(items, method, extra_body):
+            response_items.extend(self._validate_id_response(response))
+        return response_items
+
+    def _chunk_requests(
+        self,
+        items: Sequence[_T_BaseModel],
+        method: APIMethod,
+        extra_body: dict[str, Any] | None = None,
+    ) -> Iterable[SuccessResponse]:
+        """Send requests in chunks and yield responses.
+
+        Args:
+            items: The items to process.
+            method: The API method to use. This is used ot look the up the endpoint.
+            extra_body: Optional extra body content to include in the request.
+
+        Yields:
+            The successful responses from the API.
+        """
+        # Filter out None
+        endpoint = self._method_endpoint_map[method]
+
+        for chunk in chunker_sequence(items, endpoint.item_limit):
+            request = SimpleBodyRequest(
+                endpoint_url=self._make_url(endpoint.path),
+                method=endpoint.method,
+                body=TypeAdapter(dict[str, JsonValue]).dump_json(
+                    {
+                        "items": [item.model_dump(by_alias=True) for item in chunk],
+                        **(extra_body or {}),
+                    },
+                ),
+            )
+            response = self._http_client.request_with_retries(request)
+            response.raise_for_status()
+            yield response.success_response
+
+    @classmethod
+    def _filter_out_none_values(cls, params: dict[str, Any] | None) -> dict[str, Any] | None:
+        request_params: dict[str, Any] | None = None
+        if params:
+            request_params = {k: v for k, v in params.items() if v is not None}
+        return request_params
+
+    def _paginate(
+        self,
+        limit: int,
+        cursor: str | None = None,
+        params: dict[str, Any] | None = None,
+    ) -> PagedResponse[T_Response]:
+        """Fetch a single page of resources.
+
+        Args:
+            params: Query parameters for the request. Supported parameters depend on
+                the resource type but typically include:
+                - cursor: Cursor for pagination
+                - limit: Maximum number of items (defaults to list limit)
+                - space: Filter by space
+                - includeGlobal: Whether to include global resources
+            limit: Maximum number of items to return in the page.
+            cursor: Cursor for pagination.
+
+        Returns:
+            A Page containing the items and the cursor for the next page.
+        """
+        endpoint = self._method_endpoint_map["list"]
+        if not (0 < limit <= endpoint.item_limit):
+            raise ValueError(f"Limit must be between 1 and {endpoint.item_limit}, got {limit}.")
+        if endpoint.method != "GET":
+            raise NotImplementedError(f"Pagination not implemented for method {endpoint.method}.")
+        request_params = self._filter_out_none_values(params) or {}
+        request_params["limit"] = limit
+        if cursor is not None:
+            request_params["cursor"] = cursor
+        request = ParametersRequest(
+            endpoint_url=self._make_url(endpoint.path),
+            method=endpoint.method,
+            parameters=request_params,
+        )
+        result = self._http_client.request_with_retries(request)
+        result.raise_for_status()
+        return self._validate_page_response(result.success_response)
+
+    def _iterate(
+        self,
+        limit: int | None = None,
+        cursor: str | None = None,
+        params: dict[str, Any] | None = None,
+    ) -> Iterable[list[T_Response]]:
+        """Iterate over all resources, handling pagination automatically."""
+        next_cursor = cursor
+        total = 0
+        endpoint = self._method_endpoint_map["list"]
+        while True:
+            page_limit = endpoint.item_limit if limit is None else min(limit - total, endpoint.item_limit)
+            page = self._paginate(limit=page_limit, cursor=next_cursor, params=params)
+            yield page.items
+            total += len(page.items)
+            if page.next_cursor is None or (limit is not None and total >= limit):
+                break
+            next_cursor = page.next_cursor
+
+    def _list(
+        self,
+        limit: int | None = None,
+        params: dict[str, Any] | None = None,
+    ) -> list[T_Response]:
+        """List all resources, handling pagination automatically."""
+        return [item for batch in self._iterate(limit=limit, params=params) for item in batch]

cognite_neat-1.0.26/cognite/neat/_client/containers_api.py

@@ -0,0 +1,77 @@
+from collections.abc import Sequence
+
+from cognite.neat._data_model.models.dms import ContainerReference, ContainerRequest, ContainerResponse
+from cognite.neat._utils.http_client import HTTPClient, SuccessResponse
+
+from .api import Endpoint, NeatAPI
+from .config import NeatClientConfig
+from .data_classes import PagedResponse
+from .filters import ContainerFilter
+
+
+class ContainersAPI(NeatAPI):
+    def __init__(self, neat_config: NeatClientConfig, http_client: HTTPClient) -> None:
+        super().__init__(
+            neat_config,
+            http_client,
+            endpoint_map={
+                "apply": Endpoint("POST", "models/containers", item_limit=100),
+                "retrieve": Endpoint("POST", "models/containers/byids", item_limit=100),
+                "delete": Endpoint("POST", "models/containers/delete", item_limit=100),
+                "list": Endpoint("GET", "models/containers", item_limit=1000),
+            },
+        )
+
+    def _validate_page_response(self, response: SuccessResponse) -> PagedResponse[ContainerResponse]:
+        return PagedResponse[ContainerResponse].model_validate_json(response.body)
+
+    def _validate_id_response(self, response: SuccessResponse) -> list[ContainerReference]:
+        return PagedResponse[ContainerReference].model_validate_json(response.body).items
+
+    def apply(self, items: Sequence[ContainerRequest]) -> list[ContainerResponse]:
+        """Apply (create or update) containers in CDF.
+
+        Args:
+            items: List of ContainerRequest objects to apply.
+        Returns:
+            List of ContainerResponse objects.
+        """
+        return self._request_item_response(items, "apply")
+
+    def retrieve(self, items: list[ContainerReference]) -> list[ContainerResponse]:
+        """Retrieve containers by their identifiers.
+
+        Args:
+            items: List of (space, external_id) tuples identifying the containers to retrieve.
+
+        Returns:
+            List of ContainerResponse objects.
+        """
+        return self._request_item_response(items, "retrieve")
+
+    def delete(self, items: list[ContainerReference]) -> list[ContainerReference]:
+        """Delete containers by their identifiers.
+
+        Args:
+            items: List of ContainerReference objects identifying the containers to delete.
+
+        Returns:
+            List of ContainerReference objects representing the deleted containers.
+        """
+        return self._request_id_response(items, "delete")
+
+    def list(
+        self, space: str | None = None, include_global: bool = False, limit: int | None = 10
+    ) -> list[ContainerResponse]:
+        """List containers in CDF Project.
+
+        Args:
+            space: If specified, only containers in this space are returned.
+            include_global: If True, include global containers.
+            limit: Maximum number of containers to return. If None, return all containers.
+
+        Returns:
+            List of ContainerResponse objects.
+        """
+        filter = ContainerFilter(space=space, include_global=include_global)
+        return self._list(limit=limit, params=filter.dump())

cognite_neat-1.0.26/cognite/neat/_client/data_model_api.py

@@ -0,0 +1,87 @@
+from collections.abc import Sequence
+
+from cognite.neat._data_model.models.dms import DataModelReference, DataModelRequest, DataModelResponse
+from cognite.neat._utils.http_client import HTTPClient, SuccessResponse
+
+from .api import Endpoint, NeatAPI
+from .config import NeatClientConfig
+from .data_classes import PagedResponse
+from .filters import DataModelFilter
+
+
+class DataModelsAPI(NeatAPI):
+    def __init__(self, neat_config: NeatClientConfig, http_client: HTTPClient) -> None:
+        super().__init__(
+            neat_config,
+            http_client,
+            endpoint_map={
+                "apply": Endpoint("POST", "/models/datamodels", item_limit=100),
+                "retrieve": Endpoint("POST", "/models/datamodels/byids", item_limit=100),
+                "delete": Endpoint("POST", "/models/datamodels/delete", item_limit=100),
+                "list": Endpoint("GET", "/models/datamodels", item_limit=1000),
+            },
+        )
+
+    def _validate_page_response(self, response: SuccessResponse) -> PagedResponse[DataModelResponse]:
+        return PagedResponse[DataModelResponse].model_validate_json(response.body)
+
+    def _validate_id_response(self, response: SuccessResponse) -> list[DataModelReference]:
+        return PagedResponse[DataModelReference].model_validate_json(response.body).items
+
+    def apply(self, data_models: Sequence[DataModelRequest]) -> list[DataModelResponse]:
+        """Apply (create or update) data models in CDF.
+
+        Args:
+            data_models: List of DataModelRequest objects to apply.
+        Returns:
+            List of DataModelResponse objects.
+        """
+        return self._request_item_response(data_models, "apply")
+
+    def retrieve(self, items: list[DataModelReference]) -> list[DataModelResponse]:
+        """Retrieve data models by their identifiers.
+
+        Args:
+            items: List of data models references identifying the data models to retrieve.
+        Returns:
+            List of DataModelResponse objects.
+        """
+        return self._request_item_response(items, "retrieve")
+
+    def delete(self, items: list[DataModelReference]) -> list[DataModelReference]:
+        """Delete data models by their identifiers.
+
+        Args:
+            items: List of data model references identifying the data models to delete.
+        Returns:
+            List of DataModelReference objects representing the deleted data models.
+        """
+        return self._request_id_response(items, "delete")
+
+    def list(
+        self,
+        space: str | None = None,
+        all_versions: bool = False,
+        inline_views: bool = False,
+        include_global: bool = False,
+        limit: int | None = 10,
+    ) -> list[DataModelResponse]:
+        """List data models in CDF Project.
+
+        Args:
+            space: If specified, only data models in this space are returned.
+            all_versions: If True, return all versions. If False, only return the latest version.
+            inline_views: If True, include views inline in the response.
+            include_global: If True, include global data models.
+            limit: Maximum number of data models to return. If None, return all data models.
+
+        Returns:
+            List of DataModelResponse objects.
+        """
+        filter = DataModelFilter(
+            space=space,
+            all_versions=all_versions,
+            inline_views=inline_views,
+            include_global=include_global,
+        )
+        return self._list(limit=limit, params=filter.dump())

cognite_neat-1.0.26/cognite/neat/_client/filters.py

@@ -0,0 +1,40 @@
+from typing import Any, Literal
+
+from pydantic import BaseModel, ConfigDict
+from pydantic.alias_generators import to_camel
+
+
+class BaseModelRequest(BaseModel):
+    """Base class for all object. This includes resources and nested objects."""
+
+
+class Filter(BaseModel):
+    model_config = ConfigDict(alias_generator=to_camel, extra="ignore", populate_by_name=True)
+
+    def dump(self, camel_case: bool = True) -> dict[str, Any]:
+        """Dump the object to a dictionary.
+
+        Args:
+            camel_case (bool): Whether to use camelCase for the keys. Default is True.
+
+        """
+        return self.model_dump(mode="json", by_alias=camel_case, exclude_unset=True)
+
+
+class DataModelingFilter(Filter):
+    space: str | None = None
+    include_global: bool | None = None
+
+
+class ContainerFilter(DataModelingFilter):
+    used_for: Literal["node", "edge", "record", "all"] | None = None
+
+
+class ViewFilter(DataModelingFilter):
+    include_inherited_properties: bool | None = None
+    all_versions: bool | None = None
+
+
+class DataModelFilter(DataModelingFilter):
+    inline_views: bool | None = None
+    all_versions: bool | None = None

{cognite_neat-1.0.24 → cognite_neat-1.0.26}/cognite/neat/_client/init/credentials.py

@@ -12,7 +12,6 @@ def get_credentials(env_vars: ClientEnvironmentVariables) -> CredentialProvider:
         "client_credentials": create_client_credentials,
         "interactive": create_interactive_credentials,
         "token": create_token_credentials,
-        "infer": create_infer_credentials,
     }
     return options[env_vars.LOGIN_FLOW](env_vars)
 
@@ -59,12 +58,3 @@ def create_token_credentials(env_vars: ClientEnvironmentVariables) -> Credential
     if not env_vars.CDF_TOKEN:
         raise ValueError("CDF_TOKEN environment variable must be set for token authentication.")
     return Token(env_vars.CDF_TOKEN)
-
-
-def create_infer_credentials(env_vars: ClientEnvironmentVariables) -> CredentialProvider:
-    if env_vars.IDP_CLIENT_SECRET:
-        return create_client_credentials(env_vars)
-    elif env_vars.CDF_TOKEN:
-        return create_token_credentials(env_vars)
-    else:
-        return create_interactive_credentials(env_vars)

{cognite_neat-1.0.24 → cognite_neat-1.0.26}/cognite/neat/_client/init/env_vars.py

@@ -4,7 +4,6 @@ from typing import Any, Literal, TypeAlias, get_args
 
 from pydantic import BaseModel, ConfigDict, ValidationError
 
-from cognite.neat._utils.repo import get_repo_root
 from cognite.neat._utils.validation import humanize_validation_error
 
 if sys.version_info >= (3, 11):
@@ -12,9 +11,10 @@ if sys.version_info >= (3, 11):
 else:
     from typing_extensions import Self
 
-LoginFlow: TypeAlias = Literal["infer", "client_credentials", "interactive", "token"]
-VALID_LOGIN_FLOWS = get_args(LoginFlow)
+LoginFlow: TypeAlias = Literal["client_credentials", "interactive", "token"]
+AVAILABLE_LOGIN_FLOWS: tuple[LoginFlow, ...] = get_args(LoginFlow)
 Provider: TypeAlias = Literal["entra_id", "auth0", "cdf", "other"]
+AVAILABLE_PROVIDERS: tuple[Provider, ...] = get_args(Provider)
 
 
 class ClientEnvironmentVariables(BaseModel):
@@ -24,7 +24,7 @@ class ClientEnvironmentVariables(BaseModel):
     CDF_CLUSTER: str
     CDF_PROJECT: str
     PROVIDER: Provider = "entra_id"
-    LOGIN_FLOW: LoginFlow = "infer"
+    LOGIN_FLOW: LoginFlow = "client_credentials"
 
     IDP_CLIENT_ID: str | None = None
     IDP_CLIENT_SECRET: str | None = None
@@ -103,24 +103,7 @@ class ClientEnvironmentVariables(BaseModel):
         )
 
 
-def get_environment_variables(env_file_name: str) -> ClientEnvironmentVariables:
-    to_search: list[tuple[str, Path]] = []
-    try:
-        repo_root = get_repo_root()
-    except RuntimeError:
-        ...
-    else:
-        to_search.append(("repository root", repo_root))
-    to_search.append(("current working directory", Path.cwd()))
-    for location_desc, path in to_search:
-        env_path = path / env_file_name
-        if env_path.is_file():
-            print(f"Found {env_file_name} in {location_desc}.")
-            return _parse_env_file(env_path)
-    raise FileNotFoundError(f"Could not find {env_file_name} in the repository root or current working directory.")
-
-
-def _parse_env_file(env_file_path: Path) -> ClientEnvironmentVariables:
+def parse_env_file(env_file_path: Path) -> ClientEnvironmentVariables:
     content = env_file_path.read_text()
     variables: dict[str, Any] = {}
     for line in content.splitlines():
@@ -129,3 +112,32 @@ def _parse_env_file(env_file_path: Path) -> ClientEnvironmentVariables:
         key, value = line.strip().split("=", 1)
         variables[key] = value
     return ClientEnvironmentVariables.create_humanize(variables)
+
+
+def create_env_file_content(provider: Provider, login_flow: LoginFlow) -> str:
+    lines = [
+        "# Cognite NEAT Client Environment Variables",
+        "CDF_CLUSTER=<your-cdf-cluster>",
+        "CDF_PROJECT=<your-cdf-project>",
+        "",
+    ]
+    if login_flow != "token":
+        lines.append(f"PROVIDER={provider}")
+    lines.append(f"LOGIN_FLOW={login_flow}")
+    lines.append("")
+    if login_flow in ("client_credentials", "interactive"):
+        lines.append("IDP_CLIENT_ID=<your-idp-client-id>")
+        if login_flow == "client_credentials":
+            lines.append("IDP_CLIENT_SECRET=<your-idp-client-secret>")
+        if provider == "entra_id":
+            lines.append("IDP_TENANT_ID=<your-idp-tenant-id>")
+        if provider not in ("cdf", "entra_id"):
+            lines.append("IDP_TOKEN_URL=<your-idp-token-url>")
+        if provider == "other":
+            lines.append("IDP_AUDIENCE=<your-idp-audience>")
+            lines.append("IDP_SCOPES=<your-idp-scopes-comma-separated>")
+            lines.append("IDP_AUTHORITY_URL=<your-idp-authority-url>")
+    elif login_flow == "token":
+        lines.append("CDF_TOKEN=<your-cdf-token>")
+
+    return "\n".join(lines)

cognite_neat-1.0.26/cognite/neat/_client/init/interactive.py

@@ -0,0 +1,133 @@
+import importlib.util
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any
+
+from .env_vars import AVAILABLE_LOGIN_FLOWS, AVAILABLE_PROVIDERS, LoginFlow, Provider, create_env_file_content
+
+
+class InteractiveFlow(ABC):
+    def __init__(self, env_path: Path):
+        self.env_path = env_path
+
+    @abstractmethod
+    def run(self) -> None: ...
+
+    def create_env_file(self, provider: Provider, login_flow: LoginFlow) -> None:
+        env_content = create_env_file_content(provider, login_flow)
+        self.env_path.write_text(env_content, encoding="utf-8", newline="\n")
+
+
+def get_interactive_flow(env_file_path: Path) -> InteractiveFlow:
+    try:
+        importlib.util.find_spec("IPython")
+        importlib.util.find_spec("ipywidgets")
+        if not _is_in_notebook():
+            return NoDependencyFlow(env_file_path)
+        return NotebookFlow(env_file_path)
+    except ImportError:
+        return NoDependencyFlow(env_file_path)
+
+
+def _is_in_notebook() -> bool:
+    try:
+        from IPython import get_ipython
+
+        if "IPKernelApp" not in get_ipython().config:
+            return False
+    except ImportError:
+        return False
+    except AttributeError:
+        return False
+    return True
+
+
+class NoDependencyFlow(InteractiveFlow):
+    def run(self) -> None:
+        if not self.should_create_env_file():
+            return None
+        provider = self.provider()
+        login_flow: LoginFlow
+        if provider != "cdf":
+            login_flow = self.login_flow()
+        else:
+            login_flow = "client_credentials"
+        self.create_env_file(provider, login_flow)
+        print(f"Created environment file at {self.env_path!r}.")
+        return None
+
+    def should_create_env_file(self) -> bool:
+        env_file_name = self.env_path.name
+        answer = input(
+            f"Would you like to create a new {env_file_name!r} file with the required environment variables? [y/N]: "
+        )
+        return answer.strip().lower() == "y"
+
+    @classmethod
+    def provider(cls) -> Provider:
+        index = cls._prompt_choice(AVAILABLE_PROVIDERS, "Select provider:")
+        return AVAILABLE_PROVIDERS[index]
+
+    @classmethod
+    def login_flow(cls) -> LoginFlow:
+        index = cls._prompt_choice(AVAILABLE_LOGIN_FLOWS, "Select login flow:")
+        return AVAILABLE_LOGIN_FLOWS[index]
+
+    @classmethod
+    def _prompt_choice(cls, options: tuple[str, ...], prompt: str) -> int:
+        for i, option in enumerate(options, start=1):
+            print(f"{i}. {option}")
+        question = f"{prompt} [1-{len(options)}]: "
+        while True:
+            answer = input(question)
+            if answer.isdigit():
+                index = int(answer)
+                if 1 <= index <= len(options):
+                    return index - 1
+            print(f"Invalid input. Please enter a number between 1 and {len(options)}.")
+
+
+class NotebookFlow(InteractiveFlow):
+    def __init__(self, env_path: Path):
+        super().__init__(env_path)
+        import ipywidgets as widgets  # type: ignore[import-untyped]
+        from IPython.display import display
+
+        self._widgets = widgets
+        self._display = display
+
+    def run(self) -> None:
+        dropdown_providers = self._widgets.Dropdown(
+            options=list(AVAILABLE_PROVIDERS),
+            value=AVAILABLE_PROVIDERS[0],
+            description="Provider:",
+        )
+        dropdown_login_flows = self._widgets.Dropdown(
+            options=list(AVAILABLE_LOGIN_FLOWS),
+            value=AVAILABLE_LOGIN_FLOWS[0],
+            description="Login Flow:",
+        )
+        confirm_button = self._widgets.Button(description="Create template .env file", button_style="primary")
+        dropdowns = self._widgets.HBox([dropdown_providers, dropdown_login_flows])
+        output = self._widgets.Output()
+        container = self._widgets.VBox([dropdowns, confirm_button, output])
+
+        self._display(container)
+
+        def on_confirm_clicked(b: Any) -> None:
+            with output:
+                provider = dropdown_providers.value
+                login_flow = dropdown_login_flows.value
+                if provider == "cdf" and login_flow != "client_credentials":
+                    print(
+                        "Warning: 'cdf' provider only supports 'client_credentials' login flow. Overriding selection."
+                    )
+                    login_flow = "client_credentials"
+                is_existing = self.env_path.exists()
+                self.create_env_file(provider, login_flow)
+                if is_existing:
+                    print(f"Overwrote existing environment file at {self.env_path!r}.")
+                else:
+                    print(f"Created environment file at {self.env_path!r}.")
+
+        confirm_button.on_click(on_confirm_clicked)