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

cognite/neat/__init__.py

@@ -1,6 +1,7 @@
-from cognite.neat.core._utils.auth import get_cognite_client
+from cognite.neat._v0.core._utils.auth import get_cognite_client
 
+from ._config import NeatConfig, get_neat_config_from_file
+from ._session import NeatSession
 from ._version import __version__
-from .session import NeatSession
 
-__all__ = ["NeatSession", "__version__", "get_cognite_client"]
+__all__ = ["NeatConfig", "NeatSession", "__version__", "get_cognite_client", "get_neat_config_from_file"]

cognite/neat/_client/__init__.py

@@ -0,0 +1,5 @@
+from .client import NeatClient
+from .config import NeatClientConfig
+from .init.main import get_cognite_client
+
+__all__ = ["NeatClient", "NeatClientConfig", "get_cognite_client"]

cognite/neat/_client/api.py

@@ -0,0 +1,8 @@
+from cognite.neat._client.config import NeatClientConfig
+from cognite.neat._utils.http_client import HTTPClient
+
+
+class NeatAPI:
+    def __init__(self, neat_config: NeatClientConfig, http_client: HTTPClient) -> None:
+        self._config = neat_config
+        self._http_client = http_client

cognite/neat/_client/client.py

@@ -0,0 +1,21 @@
+from cognite.client import ClientConfig, CogniteClient
+
+from cognite.neat._utils.http_client import HTTPClient
+
+from .config import NeatClientConfig
+from .containers_api import ContainersAPI
+from .data_model_api import DataModelsAPI
+from .spaces_api import SpacesAPI
+from .statistics_api import StatisticsAPI
+from .views_api import ViewsAPI
+
+
+class NeatClient:
+    def __init__(self, cognite_client_or_config: CogniteClient | ClientConfig) -> None:
+        self.config = NeatClientConfig(cognite_client_or_config)
+        self.http_client = HTTPClient(self.config)
+        self.data_models = DataModelsAPI(self.config, self.http_client)
+        self.views = ViewsAPI(self.config, self.http_client)
+        self.containers = ContainersAPI(self.config, self.http_client)
+        self.spaces = SpacesAPI(self.config, self.http_client)
+        self.statistics = StatisticsAPI(self.config, self.http_client)

cognite/neat/_client/config.py

@@ -0,0 +1,40 @@
+from cognite.client import ClientConfig, CogniteClient
+
+
+class NeatClientConfig(ClientConfig):
+    def __init__(self, cognite_client_or_config: CogniteClient | ClientConfig):
+        config = (
+            cognite_client_or_config.config
+            if isinstance(cognite_client_or_config, CogniteClient)
+            else cognite_client_or_config
+        )
+        super().__init__(
+            client_name=config.client_name,
+            project=config.project,
+            credentials=config.credentials,
+            api_subversion=config.api_subversion,
+            base_url=config.base_url,
+            max_workers=config.max_workers,
+            headers=config.headers,
+            timeout=config.timeout,
+            file_transfer_timeout=config.file_transfer_timeout,
+            debug=config.debug,
+        )
+
+    def create_api_url(self, endpoint: str) -> str:
+        """Create a full API URL for the given endpoint.
+
+        Args:
+            endpoint (str): The API endpoint to append to the base URL.
+
+        Returns:
+            str: The full API URL.
+
+        Examples:
+            >>> config = NeatClientConfig(cluster="bluefield", project="my_project", ...)
+            >>> config.create_api_url("/models/instances")
+            "https://bluefield.cognitedata.com/api/v1/my_project/models/instances"
+        """
+        if not endpoint.startswith("/"):
+            endpoint = f"/{endpoint}"
+        return f"{self.base_url}/api/v1/projects/{self.project}{endpoint}"

cognite/neat/_client/containers_api.py

@@ -0,0 +1,138 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from cognite.neat._data_model.models.dms import ContainerReference, ContainerRequest, ContainerResponse, DataModelBody
+from cognite.neat._utils.http_client import ItemIDBody, ItemsRequest, ParametersRequest
+from cognite.neat._utils.useful_types import PrimitiveType
+
+from .api import NeatAPI
+from .data_classes import PagedResponse
+
+
+class ContainersAPI(NeatAPI):
+    ENDPOINT = "/models/containers"
+    LIST_REQUEST_LIMIT = 1000
+
+    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.
+        """
+        if not items:
+            return []
+        if len(items) > 100:
+            raise ValueError("Cannot apply more than 100 containers at once.")
+        result = self._http_client.request_with_retries(
+            ItemsRequest(
+                endpoint_url=self._config.create_api_url(self.ENDPOINT),
+                method="POST",
+                body=DataModelBody(items=items),
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[ContainerResponse].model_validate_json(result.success_response.body)
+        return result.items
+
+    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.
+        """
+        if not items:
+            return []
+        if len(items) > 1000:
+            raise ValueError("Cannot retrieve more than 1000 containers at once.")
+
+        result = self._http_client.request_with_retries(
+            ItemsRequest(
+                endpoint_url=self._config.create_api_url(f"{self.ENDPOINT}/byids"),
+                method="POST",
+                body=ItemIDBody(items=items),
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[ContainerResponse].model_validate_json(result.success_response.body)
+        return result.items
+
+    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.
+        """
+        if not items:
+            return []
+        if len(items) > 100:
+            raise ValueError("Cannot delete more than 100 containers at once.")
+
+        result = self._http_client.request_with_retries(
+            ItemsRequest(
+                endpoint_url=self._config.create_api_url(f"{self.ENDPOINT}/delete"),
+                method="POST",
+                body=ItemIDBody(items=items),
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[ContainerReference].model_validate_json(result.success_response.body)
+        return result.items
+
+    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.
+        """
+        if limit is not None and limit < 0:
+            raise ValueError("Limit must be non-negative.")
+        elif limit is not None and limit == 0:
+            return []
+        parameters: dict[str, PrimitiveType] = {"includeGlobal": include_global}
+        if space is not None:
+            parameters["space"] = space
+        cursor: str | None = None
+        container_responses: list[ContainerResponse] = []
+        while True:
+            if cursor is not None:
+                parameters["cursor"] = cursor
+            if limit is None:
+                parameters["limit"] = self.LIST_REQUEST_LIMIT
+            else:
+                parameters["limit"] = min(self.LIST_REQUEST_LIMIT, limit - len(container_responses))
+            result = self._http_client.request_with_retries(
+                ParametersRequest(
+                    endpoint_url=self._config.create_api_url(self.ENDPOINT),
+                    method="GET",
+                    parameters=parameters,
+                )
+            )
+            result.raise_for_status()
+            result = PagedResponse[ContainerResponse].model_validate_json(result.success_response.body)
+            container_responses.extend(result.items)
+            cursor = result.next_cursor
+            if cursor is None or (limit is not None and len(container_responses) >= limit):
+                break
+        return container_responses

cognite/neat/_client/data_classes.py

@@ -0,0 +1,44 @@
+from typing import Generic, TypeVar
+
+from pydantic import BaseModel, Field
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class PagedResponse(BaseModel, Generic[T]):
+    items: list[T]
+    next_cursor: str | None = Field(None, alias="nextCursor")
+
+
+class ResourceLimit(BaseModel):
+    """Model for resources with count and limit."""
+
+    count: int
+    limit: int
+
+
+class InstancesDetail(BaseModel, populate_by_name=True):
+    """Model for instances with detailed metrics."""
+
+    edges: int
+    soft_deleted_edges: int = Field(alias="softDeletedEdges")
+    nodes: int
+    soft_deleted_nodes: int = Field(alias="softDeletedNodes")
+    instances: int
+    instances_limit: int = Field(alias="instancesLimit")
+    soft_deleted_instances: int = Field(alias="softDeletedInstances")
+    soft_deleted_instances_limit: int = Field(alias="softDeletedInstancesLimit")
+
+
+class StatisticsResponse(BaseModel, populate_by_name=True):
+    """Main API response model."""
+
+    spaces: ResourceLimit
+    containers: ResourceLimit
+    views: ResourceLimit
+    data_models: ResourceLimit = Field(alias="dataModels")
+    container_properties: ResourceLimit = Field(alias="containerProperties")
+    instances: InstancesDetail
+    concurrent_read_limit: int = Field(alias="concurrentReadLimit")
+    concurrent_write_limit: int = Field(alias="concurrentWriteLimit")
+    concurrent_delete_limit: int = Field(alias="concurrentDeleteLimit")

cognite/neat/_client/data_model_api.py

@@ -0,0 +1,115 @@
+from collections.abc import Sequence
+
+from cognite.neat._data_model.models.dms import DataModelBody, DataModelReference, DataModelRequest, DataModelResponse
+from cognite.neat._utils.http_client import ItemIDBody, ItemsRequest, ParametersRequest
+from cognite.neat._utils.useful_types import PrimitiveType
+
+from .api import NeatAPI
+from .data_classes import PagedResponse
+
+
+class DataModelsAPI(NeatAPI):
+    ENDPOINT = "/models/datamodels"
+
+    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.
+        """
+        if not data_models:
+            return []
+        if len(data_models) > 100:
+            raise ValueError("Cannot apply more than 100 data models at once.")
+
+        result = self._http_client.request_with_retries(
+            ItemsRequest(
+                endpoint_url=self._config.create_api_url(self.ENDPOINT),
+                method="POST",
+                body=DataModelBody(items=data_models),
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[DataModelResponse].model_validate_json(result.success_response.body)
+        return result.items
+
+    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.
+        """
+        if not items:
+            return []
+        if len(items) > 100:
+            raise ValueError("Cannot retrieve more than 1000 containers at once.")
+
+        result = self._http_client.request_with_retries(
+            ItemsRequest(
+                endpoint_url=self._config.create_api_url(f"{self.ENDPOINT}/byids"),
+                method="POST",
+                body=ItemIDBody(items=items),
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[DataModelResponse].model_validate_json(result.success_response.body)
+        return result.items
+
+    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.
+        """
+        if not items:
+            return []
+        if len(items) > 100:
+            raise ValueError("Cannot delete more than 100 data models at once.")
+
+        result = self._http_client.request_with_retries(
+            ItemsRequest(
+                endpoint_url=self._config.create_api_url(f"{self.ENDPOINT}/delete"),
+                method="POST",
+                body=ItemIDBody(items=items),
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[DataModelReference].model_validate_json(result.success_response.body)
+        return result.items
+
+    def list(
+        self,
+        space: str | None = None,
+        all_versions: bool = False,
+        include_global: bool = False,
+        limit: int = 10,
+    ) -> list[DataModelResponse]:
+        """List data models in CDF Project."""
+        if limit > 1000:
+            raise ValueError("Pagination is not (yet) supported for listing data models. The maximum limit is 1000.")
+        parameters: dict[str, PrimitiveType] = {
+            "allVersions": all_versions,
+            "includeGlobal": include_global,
+            "limit": limit,
+        }
+        if space is not None:
+            parameters["space"] = space
+        result = self._http_client.request_with_retries(
+            ParametersRequest(
+                endpoint_url=self._config.create_api_url(self.ENDPOINT),
+                method="GET",
+                parameters=parameters,
+            )
+        )
+        result.raise_for_status()
+        result = PagedResponse[DataModelResponse].model_validate_json(result.success_response.body)
+        return result.items

cognite/neat/_client/init/credentials.py

@@ -0,0 +1,70 @@
+from collections.abc import Callable
+
+from cognite.client.credentials import CredentialProvider, OAuthClientCredentials, OAuthInteractive, Token
+
+from cognite.neat._utils.text import humanize_collection
+
+from .env_vars import ClientEnvironmentVariables, LoginFlow
+
+
+def get_credentials(env_vars: ClientEnvironmentVariables) -> CredentialProvider:
+    options: dict[LoginFlow, Callable[[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)
+
+
+def create_client_credentials(env_vars: ClientEnvironmentVariables) -> CredentialProvider:
+    missing: list[str] = []
+    if not env_vars.IDP_CLIENT_ID:
+        missing.append("IDP_CLIENT_ID")
+    if not env_vars.IDP_CLIENT_SECRET:
+        missing.append("IDP_CLIENT_SECRET")
+    if env_vars.IDP_CLIENT_ID is None or env_vars.IDP_CLIENT_SECRET is None:
+        raise ValueError(
+            f"The following environment variables must be set for "
+            f"client credentials authentication: {humanize_collection(missing)}"
+        )
+
+    if env_vars.PROVIDER == "cdf":
+        return OAuthClientCredentials(
+            client_id=env_vars.IDP_CLIENT_ID,
+            client_secret=env_vars.IDP_CLIENT_SECRET,
+            token_url=env_vars.idp_token_url,
+            scopes=None,  # type: ignore[arg-type]
+        )
+    return OAuthClientCredentials(
+        client_id=env_vars.IDP_CLIENT_ID,
+        client_secret=env_vars.IDP_CLIENT_SECRET,
+        token_url=env_vars.idp_token_url,
+        audience=env_vars.idp_audience,
+        scopes=env_vars.idp_scopes,
+    )
+
+
+def create_interactive_credentials(env_vars: ClientEnvironmentVariables) -> CredentialProvider:
+    if not env_vars.IDP_CLIENT_ID:
+        raise ValueError("IDP_CLIENT_ID environment variable must be set for interactive authentication.")
+    return OAuthInteractive(
+        client_id=env_vars.IDP_CLIENT_ID,
+        authority_url=env_vars.idp_authority_url,
+        scopes=env_vars.idp_scopes,
+    )
+
+
+def create_token_credentials(env_vars: ClientEnvironmentVariables) -> CredentialProvider:
+    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/_client/init/env_vars.py

@@ -0,0 +1,131 @@
+import sys
+from pathlib import Path
+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):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
+LoginFlow: TypeAlias = Literal["infer", "client_credentials", "interactive", "token"]
+VALID_LOGIN_FLOWS = get_args(LoginFlow)
+Provider: TypeAlias = Literal["entra_id", "auth0", "cdf", "other"]
+
+
+class ClientEnvironmentVariables(BaseModel):
+    """Configuration for environment variables used by the NEAT client."""
+
+    model_config = ConfigDict(extra="forbid", frozen=True)
+    CDF_CLUSTER: str
+    CDF_PROJECT: str
+    PROVIDER: Provider = "entra_id"
+    LOGIN_FLOW: LoginFlow = "infer"
+
+    IDP_CLIENT_ID: str | None = None
+    IDP_CLIENT_SECRET: str | None = None
+    CDF_TOKEN: str | None = None
+
+    IDP_TENANT_ID: str | None = None
+    IDP_TOKEN_URL: str | None = None
+
+    CDF_URL: str | None = None
+    IDP_AUDIENCE: str | None = None
+    IDP_SCOPES: str | None = None
+    IDP_AUTHORITY_URL: str | None = None
+    IDP_DISCOVERY_URL: str | None = None
+    CDF_MAX_WORKERS: int | None = None
+    CDF_CLIENT_TIMEOUT: int | None = None
+    CDF_REDIRECT_PORT: int = 53_000
+
+    @classmethod
+    def create_humanize(cls, values: dict[str, Any]) -> Self:
+        try:
+            return cls.model_validate(values)
+        except ValidationError as e:
+            errors = [humanize_validation_error(error) for error in e.errors()]
+            raise ValueError("Invalid environment variable configuration:\n" + "\n - ".join(errors)) from e
+
+    @property
+    def idp_tenant_id(self) -> str:
+        if self.IDP_TENANT_ID:
+            return self.IDP_TENANT_ID
+        # This line is technically unreachable due to the checks in idp_token_url and idp_authority_url
+        raise RuntimeError("IDP_TENANT_ID is missing")
+
+    @property
+    def idp_token_url(self) -> str:
+        if self.PROVIDER == "cdf":
+            return "https://auth.cognite.com/oauth2/token"
+        if self.IDP_TOKEN_URL:
+            return self.IDP_TOKEN_URL
+        if self.PROVIDER == "entra_id" and self.IDP_TENANT_ID:
+            return f"https://login.microsoftonline.com/{self.IDP_TENANT_ID}/oauth2/v2.0/token"
+        alternative = " or provide IDP_TENANT_ID" if self.PROVIDER == "entra_id" else ""
+        raise ValueError(
+            f"IDP_TOKEN_URL is missing. Please provide it{alternative} in the environment variables.",
+        )
+
+    @property
+    def cdf_url(self) -> str:
+        return self.CDF_URL or f"https://{self.CDF_CLUSTER}.cognitedata.com"
+
+    @property
+    def idp_audience(self) -> str:
+        if self.IDP_AUDIENCE:
+            return self.IDP_AUDIENCE
+        if self.PROVIDER == "auth0":
+            return f"https://{self.CDF_PROJECT}.fusion.cognite.com/{self.CDF_PROJECT}"
+        else:
+            return f"https://{self.CDF_CLUSTER}.cognitedata.com"
+
+    @property
+    def idp_scopes(self) -> list[str]:
+        if self.IDP_SCOPES:
+            return self.IDP_SCOPES.split(",")
+        if self.PROVIDER == "auth0":
+            return ["IDENTITY", "user_impersonation"]
+        return [f"https://{self.CDF_CLUSTER}.cognitedata.com/.default"]
+
+    @property
+    def idp_authority_url(self) -> str:
+        if self.IDP_AUTHORITY_URL:
+            return self.IDP_AUTHORITY_URL
+        if self.PROVIDER == "entra_id" and self.IDP_TENANT_ID:
+            return f"https://login.microsoftonline.com/{self.IDP_TENANT_ID}"
+        alternative = " or provide IDP_TENANT_ID" if self.PROVIDER == "entra_id" else ""
+        raise ValueError(
+            f"IDP_AUTHORITY_URL is missing. Please provide it{alternative} in the environment variables.",
+        )
+
+
+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:
+    content = env_file_path.read_text()
+    variables: dict[str, Any] = {}
+    for line in content.splitlines():
+        if line.startswith("#") or "=" not in line:
+            continue
+        key, value = line.strip().split("=", 1)
+        variables[key] = value
+    return ClientEnvironmentVariables.create_humanize(variables)

cognite/neat/_client/init/main.py

@@ -0,0 +1,51 @@
+from cognite.client import CogniteClient
+from cognite.client.config import ClientConfig, global_config
+
+from cognite.neat import _version
+
+from .credentials import get_credentials
+from .env_vars import ClientEnvironmentVariables, get_environment_variables
+
+CLIENT_NAME = f"CogniteNeat:{_version.__version__}"
+
+
+def get_cognite_client(env_file_name: str) -> CogniteClient:
+    """Get a CogniteClient using environment variables from a .env file."
+
+    Args:
+        env_file_name: The name of the .env file to look for in the repository root / current working directory. If
+        the file is found, the variables will be loaded from the file. If the file is not found, the user will
+        be prompted to enter the variables and the file will be created.
+
+    Returns:
+        CogniteClient: An instance of CogniteClient configured with the loaded environment variables.
+    """
+    try:
+        return get_cognite_client_internal(env_file_name)
+    except Exception as e:
+        raise RuntimeError(f"Failed to create client ❌: {e!s}") from None
+
+
+def get_cognite_client_internal(env_file_name: str) -> CogniteClient:
+    # This function raises exceptions on failure
+    if not env_file_name.endswith(".env"):
+        raise ValueError(f"env_file_name must end with '.env'. Got: {env_file_name!r}")
+    global_config.disable_pypi_version_check = True
+    global_config.silence_feature_preview_warnings = True
+    env_vars = get_environment_variables(env_file_name)
+    client_config = create_client_config_from_env_vars(env_vars)
+    # Todo validate credentials by making a simple call to CDF
+    #   Offer to store credentials securely if valid
+    #
+    return CogniteClient(client_config)
+
+
+def create_client_config_from_env_vars(env_vars: ClientEnvironmentVariables) -> ClientConfig:
+    return ClientConfig(
+        client_name=CLIENT_NAME,
+        project=env_vars.CDF_PROJECT,
+        credentials=get_credentials(env_vars),
+        max_workers=env_vars.CDF_MAX_WORKERS,
+        timeout=env_vars.CDF_CLIENT_TIMEOUT,
+        base_url=env_vars.cdf_url,
+    )