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

cognite/neat/_session/_result/_deployment/_physical/_statistics.py

@@ -0,0 +1,180 @@
+import itertools
+from typing import Any, Literal, cast, get_args
+
+from pydantic import BaseModel, Field
+
+from cognite.neat._data_model.deployer.data_classes import (
+    AppliedChanges,
+    DataModelEndpoint,
+    DeploymentResult,
+)
+
+
+class EndpointStatistics(BaseModel):
+    """Statistics for a single endpoint."""
+
+    create: int = 0
+    update: int = 0
+    delete: int = 0
+    skip: int = 0
+    unchanged: int = 0
+    failed: int = 0
+
+    @property
+    def total(self) -> int:
+        """Total number of changes for this endpoint."""
+        return self.create + self.update + self.delete + self.skip + self.unchanged + self.failed
+
+    def increment(self, change_type: Literal["create", "update", "delete", "unchanged", "skip", "failed"]) -> None:
+        """Increment the count for a specific change type.
+
+        Args:
+            change_type: The type of change (create, update, delete, skip, unchanged, failed).
+        """
+        if hasattr(self, change_type):
+            setattr(self, change_type, getattr(self, change_type) + 1)
+        else:
+            raise RuntimeError(f"Unknown change type: {change_type}. This is a bug in NEAT.")
+
+
+class ChangeTypeStatistics(BaseModel):
+    """Statistics grouped by change type."""
+
+    create: int = 0
+    update: int = 0
+    delete: int = 0
+    skip: int = 0
+    unchanged: int = 0
+    failed: int = 0
+
+    def increment(self, change_type: Literal["create", "update", "delete", "unchanged", "skip", "failed"]) -> None:
+        """Increment the count for a specific change type.
+
+        Args:
+            change_type: The type of change (create, update, delete, skip, unchanged, failed).
+        """
+        if hasattr(self, change_type):
+            setattr(self, change_type, getattr(self, change_type) + 1)
+        else:
+            raise RuntimeError(f"Unknown change type: {change_type}. This is a bug in NEAT.")
+
+
+class SeverityStatistics(BaseModel):
+    """Statistics grouped by severity."""
+
+    SAFE: int = 0
+    WARNING: int = 0
+    BREAKING: int = 0
+
+    def increment(self, severity: Literal["SAFE", "WARNING", "BREAKING"]) -> None:
+        """Increment the count for a specific severity level.
+
+        Args:
+            severity: The severity level (SAFE, WARNING, BREAKING).
+        """
+        if hasattr(self, severity):
+            setattr(self, severity, getattr(self, severity) + 1)
+        else:
+            raise RuntimeError(f"Unknown severity level: {severity}. This is a bug in NEAT.")
+
+
+class DeploymentStatistics(BaseModel):
+    """Overall deployment statistics."""
+
+    by_endpoint: dict[str, EndpointStatistics] = Field(default_factory=dict)
+    by_change_type: ChangeTypeStatistics = Field(default_factory=ChangeTypeStatistics)
+    by_severity: SeverityStatistics = Field(default_factory=SeverityStatistics)
+
+    @property
+    def total_changes(self) -> int:
+        """Total number of changes in the deployment."""
+        return sum(endpoint_stats.total for endpoint_stats in self.by_endpoint.values())
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        """Include computed properties in serialization."""
+        data = super().model_dump(**kwargs)
+        data["total_changes"] = self.total_changes
+        return data
+
+    @classmethod
+    def from_deployment_result(cls, result: DeploymentResult) -> "DeploymentStatistics":
+        """Create DeploymentStatistics from a DeploymentResult.
+
+        Args:
+            result: The deployment result to create statistics from.
+
+        Returns:
+            DeploymentStatistics instance with computed statistics.
+        """
+        stats = cls(
+            by_endpoint={endpoint: EndpointStatistics() for endpoint in get_args(DataModelEndpoint)},
+            by_change_type=ChangeTypeStatistics(),
+            by_severity=SeverityStatistics(),
+        )
+
+        if result.is_dry_run:
+            stats._update_from_dry_run(result)
+        else:
+            stats._update_from_deployment(result)
+
+        return stats
+
+    def _update_from_dry_run(self, result: DeploymentResult) -> None:
+        """Update statistics from dry run mode.
+
+        Args:
+            result: The deployment result in dry run mode.
+        """
+        for plan in result.plan:
+            for resource in plan.resources:
+                self._update_single_stat(
+                    endpoint=plan.endpoint,
+                    change_type=resource.change_type,
+                    severity=cast(Literal["SAFE", "WARNING", "BREAKING"], resource.severity.name),
+                )
+
+    def _update_from_deployment(self, result: DeploymentResult) -> None:
+        """Update statistics from actual deployment mode.
+
+        Args:
+            result: The deployment result from actual deployment.
+        """
+        applied_changes = cast(AppliedChanges, result.responses)
+
+        for response in itertools.chain(
+            applied_changes.created,
+            applied_changes.merged_updated,
+            applied_changes.deletions,
+            applied_changes.unchanged,
+            applied_changes.skipped,
+        ):
+            self._update_single_stat(
+                endpoint=response.endpoint,
+                change_type=response.change.change_type if response.is_success else "failed",
+                severity=cast(Literal["SAFE", "WARNING", "BREAKING"], response.change.severity.name),
+            )
+
+    def _update_single_stat(
+        self,
+        endpoint: str,
+        change_type: Literal["create", "update", "delete", "unchanged", "skip", "failed"],
+        severity: Literal["SAFE", "WARNING", "BREAKING"],
+    ) -> None:
+        """Update all statistics for a single change.
+
+        Args:
+            endpoint: The endpoint type (spaces, containers, views, datamodels).
+            change_type: The type of change (create, update, delete, skip, unchanged, failed).
+            severity: The severity level (SAFE, WARNING, BREAKING).
+        """
+        # Update by change type statistics
+        self.by_change_type.increment(change_type)
+
+        # Update by endpoint statistics
+        if endpoint in self.by_endpoint:
+            self.by_endpoint[endpoint].increment(change_type)
+        else:
+            raise RuntimeError(f"Unknown endpoint: {endpoint}. This is a bug in NEAT.")
+
+        # Update severity statistics
+        self.by_severity.increment(severity)

cognite/neat/_session/_result/_deployment/_physical/serializer.py

@@ -0,0 +1,35 @@
+import uuid
+from typing import Any
+
+from cognite.neat._data_model.deployer.data_classes import DeploymentResult
+
+from ._changes import SerializedChanges
+from ._statistics import DeploymentStatistics
+
+
+def serialize_deployment_result(result: DeploymentResult) -> dict[str, Any]:
+    """Serialize deployment result into structured changes.
+
+    Args:
+        result: The deployment result to serialize.
+
+    Returns:
+        Serialized changes representing the deployment result.
+    """
+    result_dict = {"unique_id": uuid.uuid4().hex[:8], "status": result.status, "is_dry_run": result.is_dry_run}
+
+    stats = DeploymentStatistics.from_deployment_result(result)
+    changes = SerializedChanges.from_deployment_result(result)
+
+    result_dict["total_changes"] = stats.total_changes
+    result_dict["created"] = stats.by_change_type.create
+    result_dict["updated"] = stats.by_change_type.update
+    result_dict["deleted"] = stats.by_change_type.delete
+    result_dict["skipped"] = stats.by_change_type.skip
+    result_dict["unchanged"] = stats.by_change_type.unchanged
+    result_dict["failed"] = stats.by_change_type.failed
+
+    result_dict["STATS_JSON"] = stats.model_dump_json()
+    result_dict["CHANGES_JSON"] = changes.model_dump_json_flat()
+
+    return result_dict

cognite/neat/_session/_result/_result.py

@@ -0,0 +1,31 @@
+from cognite.neat._data_model.deployer.data_classes import DeploymentResult
+from cognite.neat._session._html._render import render
+from cognite.neat._store import NeatStore
+
+from ._deployment._physical.serializer import serialize_deployment_result
+
+
+class Result:
+    """Class to handle deployment results in the NeatSession."""
+
+    def __init__(self, store: NeatStore) -> None:
+        self._store = store
+
+    @property
+    def _result(self) -> DeploymentResult | None:
+        """Get deployment result from the last change in the store."""
+        if change := self._store.provenance.last_change:
+            if change.result:
+                return change.result
+        return None
+
+    def _repr_html_(self) -> str:
+        """Generate interactive HTML representation."""
+        if not self._result:
+            return "<p>No deployment result available</p>"
+
+        if isinstance(self._result, DeploymentResult):
+            serialized_result = serialize_deployment_result(self._result)
+            return render("deployment", serialized_result)
+        else:
+            raise NotImplementedError(f"HTML rendering for the result type {type(self._result)} is not implemented.")

cognite/neat/_session/_session.py

@@ -0,0 +1,81 @@
+import json
+
+from cognite.client import ClientConfig, CogniteClient
+
+from cognite.neat import _version
+from cognite.neat._client import NeatClient
+from cognite.neat._config import NeatConfig, PredefinedProfile
+from cognite.neat._session._usage_analytics._collector import Collector
+from cognite.neat._state_machine import EmptyState, PhysicalState
+from cognite.neat._store import NeatStore
+from cognite.neat._utils.http_client import ParametersRequest, SuccessResponse
+
+from ._issues import Issues
+from ._physical import PhysicalDataModel
+from ._result import Result
+
+
+class NeatSession:
+    """A session is an interface for neat operations."""
+
+    def __init__(
+        self, client: CogniteClient | ClientConfig, config: PredefinedProfile | NeatConfig = "legacy-additive"
+    ) -> None:
+        """Initialize a Neat session.
+
+        Args:
+            client (CogniteClient | ClientConfig): The Cognite client or client configuration to use for the session.
+            config (Literal["legacy-additive", "legacy-rebuild", "deep-additive", "deep-rebuild"] | NeatConfig):
+                The configuration profile to use for the session.
+                Defaults to "legacy-additive". This means Neat will perform additive modeling
+                and apply only validations that were part of the legacy Neat version.
+        """
+        self._config = NeatConfig.create_predefined(config) if isinstance(config, str) else config
+
+        # Use configuration for physical data model
+        self._client = NeatClient(client)
+        self._store = NeatStore(config=self._config, client=self._client)
+        self.physical_data_model = PhysicalDataModel(self._store, self._client, self._config)
+        self.issues = Issues(self._store)
+        self.result = Result(self._store)
+
+        collector = Collector()
+        if collector.can_collect:
+            collector.collect("initSession", {"mode": self._config.modeling.mode})
+
+        self._welcome_message()
+
+    def _welcome_message(self) -> None:
+        cdf_project = self._client.config.project
+        message = f"Neat session started for CDF project: '{cdf_project}'"
+        responses = self._client.http_client.request(
+            ParametersRequest(endpoint_url=self._client.config.create_api_url(""), method="GET")
+        )
+        if len(responses) == 1 and isinstance(response := responses[0], SuccessResponse):
+            organization = ""
+            try:
+                organization = json.loads(response.body)["organization"]
+            except (KeyError, ValueError):
+                ...
+            if organization:
+                message += f" (Organization: '{organization}')"
+
+        print(message)
+        print(self._config)
+
+    @property
+    def version(self) -> str:
+        """Get the current version of neat."""
+        return _version.__version__
+
+    def _repr_html_(self) -> str:
+        if isinstance(self._store.state, EmptyState):
+            return (
+                "<strong>Empty session</strong>. Get started by reading for example physical data model"
+                " <em>.physical_data_model.read</em>"
+            )
+
+        if isinstance(self._store.state, PhysicalState):
+            return self.physical_data_model._repr_html_()
+
+        raise RuntimeError("Unknown session state, contact support.")

cognite/neat/_session/_usage_analytics/_collector.py

@@ -0,0 +1,131 @@
+import os
+import platform
+import threading
+import uuid
+from contextlib import suppress
+from typing import Any, Literal
+
+from mixpanel import Consumer, Mixpanel  # type: ignore[import-untyped]
+
+from cognite.neat._session._usage_analytics._constants import IN_NOTEBOOK, IN_PYODIDE
+from cognite.neat._session._usage_analytics._storage import ManualReadResult, ReadResult, get_storage
+
+_NEAT_MIXPANEL_TOKEN: str = "bd630ad149d19999df3989c3a3750c4f"
+
+
+class Collector:
+    """Collects usage data and sends it to Mixpanel."""
+
+    _instance: "Collector | None" = None
+    _initialized: bool = False
+
+    def __new__(cls, *args: Any, **kwargs: Any) -> "Collector":
+        # Implementing Singleton pattern
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def __init__(self, skip_tracking: bool = False) -> None:
+        if self._initialized:
+            return
+        self.mp = Mixpanel(_NEAT_MIXPANEL_TOKEN, consumer=Consumer(api_host="api-eu.mixpanel.com"))
+        self._storage = get_storage()
+        self._opt_status_key = "neat-opt-status"
+        self._distinct_id_key = "neat-distinct-id"
+        self.skip_tracking = skip_tracking
+        self._initialized = True
+        # We trigger these reads immediately to have them ready when needed.
+        # This is necessary in a Pyodide environment were the read will not be available
+        self._opt_status_read: ReadResult | None = self._storage.read(self._opt_status_key)
+        self._distinct_id_read: ReadResult | None = self._storage.read(self._distinct_id_key)
+
+    @property
+    def _opt_status(self) -> str:
+        if self._opt_status_read is None:
+            self._opt_status_read = self._storage.read(self._opt_status_key)
+        if self._opt_status_read.is_ready:
+            return self._opt_status_read.get_data()
+        return ""  # We do not have an opt status yet.
+
+    @property
+    def _distinct_id(self) -> str | None:
+        if self._distinct_id_read is None:
+            self._distinct_id_read = self._storage.read(self._distinct_id_key)
+        if self._distinct_id_read.is_ready:
+            return self._distinct_id_read.get_data()
+        return None
+
+    def bust_opt_status(self) -> None:
+        self._storage.delete(self._opt_status_key)
+        self._opt_status_read = None
+
+    @property
+    def is_opted_out(self) -> bool:
+        return False
+
+    @property
+    def is_opted_in(self) -> bool:
+        return True
+
+    def enable(self) -> None:
+        self._storage.write(self._opt_status_key, "opted-in")
+        # Override cached property
+        self._opt_status_read = ManualReadResult("opted-in")
+
+    def disable(self) -> None:
+        self._storage.write(self._opt_status_key, "opted-out")
+        # Override cached property
+        self._opt_status_read = ManualReadResult("opted-out")
+
+    def get_distinct_id(self) -> str:
+        existing_id = self._distinct_id
+        if existing_id:
+            return existing_id
+
+        distinct_id = f"{platform.system()}-{platform.python_version()}-{uuid.uuid4()!s}"
+        self._storage.write(self._distinct_id_key, distinct_id)
+        with suppress(ConnectionError):
+            self.mp.people_set(
+                distinct_id,
+                {
+                    "$os": platform.system(),
+                    "$python_version": platform.python_version(),
+                    "$distinct_id": distinct_id,
+                    "environment": self._get_environment(),
+                },
+            )
+        self._distinct_id_read = ManualReadResult(distinct_id)
+        return distinct_id
+
+    @staticmethod
+    def _get_environment() -> Literal["pyodide", "notebook", "python"]:
+        """Get the current environment the user is running in."""
+        if IN_PYODIDE:
+            return "pyodide"
+        if IN_NOTEBOOK:
+            return "notebook"
+        return "python"
+
+    @property
+    def can_collect(self) -> bool:
+        """Check if tracking is possible."""
+        return not self.skip_tracking and self.is_opted_in and "PYTEST_CURRENT_TEST" not in os.environ
+
+    def collect(
+        self, event_name: Literal["action", "initSession", "deployment"], event_properties: dict[str, Any]
+    ) -> None:
+        distinct_id = self.get_distinct_id()
+
+        def track() -> None:
+            # If we are unable to connect to Mixpanel, we don't want to crash the program
+            with suppress(ConnectionError):
+                self.mp.track(distinct_id, event_name, event_properties)
+
+        if IN_PYODIDE:
+            # We cannot spawn threads in Pyodide
+            track()
+            return None
+
+        thread = threading.Thread(target=track, daemon=True)
+        thread.start()
+        return None

cognite/neat/_session/_usage_analytics/_constants.py

@@ -0,0 +1,23 @@
+def _is_in_notebook() -> bool:
+    try:
+        from IPython import get_ipython
+
+        if "IPKernelApp" not in get_ipython().config:  # pragma: no cover
+            return False
+    except ImportError:
+        return False
+    except AttributeError:
+        return False
+    return True
+
+
+def _is_in_browser() -> bool:
+    try:
+        from pyodide.ffi import IN_BROWSER  # type: ignore [import-not-found]
+    except ModuleNotFoundError:
+        return False
+    return IN_BROWSER
+
+
+IN_PYODIDE = _is_in_browser()
+IN_NOTEBOOK = _is_in_notebook() or IN_PYODIDE