From 1a77b0303a31067312d7ae58d4c5391124aa2e21 Mon Sep 17 00:00:00 2001 From: vikramlc Date: Mon, 7 Sep 2026 17:40:59 +0530 Subject: [PATCH] feat(integrations): Add startup/checkin endpoints + tooling fixes Adds the extractor self-registration/heartbeat protocol (startup, checkin) now that service-contracts PR #3378 removed their ifdef: internal marking, plus the retry-idempotency registration and codespell config needed to support the "checkin" identifier. Co-Authored-By: Claude Sonnet 5 --- .pre-commit-config.yaml | 6 + cognite/client/_api/integrations/__init__.py | 73 +++++++ .../client/_sync_api/integrations/__init__.py | 63 +++++- .../data_classes/integrations/__init__.py | 12 ++ .../data_classes/integrations/checkin.py | 194 ++++++++++++++++++ cognite/client/utils/_url.py | 2 + .../test_integrations/test_integrations.py | 55 +++++ tests/tests_unit/test_api_client.py | 2 + .../test_data_classes/test_integrations.py | 57 +++++ 9 files changed, 463 insertions(+), 1 deletion(-) create mode 100644 cognite/client/data_classes/integrations/checkin.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 4ebce9df61..150e618a96 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -71,5 +71,11 @@ repos: hooks: - id: codespell files: ^cognite/.* + # "checkin" matches the Integrations API's actual operationId (integration_checkin) and + # URL path (/integrations/checkin), so it can't be reworded to "check-in" without diverging + # from the real API surface. Passed as a hook arg (not pyproject.toml's [tool.codespell]) + # since reading TOML config requires tomli/tomllib, which isn't guaranteed to be available + # in the hook's isolated environment on Python <3.11. args: + - --ignore-words-list=checkin - --write-changes diff --git a/cognite/client/_api/integrations/__init__.py b/cognite/client/_api/integrations/__init__.py index 75a506555c..709bd22364 100644 --- a/cognite/client/_api/integrations/__init__.py +++ b/cognite/client/_api/integrations/__init__.py @@ -9,6 +9,7 @@ from cognite.client._api.integrations.tasks import IntegrationTasksAPI from cognite.client._api_client import APIClient from cognite.client._constants import DEFAULT_LIMIT_READ +from cognite.client.data_classes.integrations.checkin import CheckinRequest, CheckinResponse, StartupRequest from cognite.client.data_classes.integrations.integrations import ( Integration, IntegrationList, @@ -236,3 +237,75 @@ async def delete(self, external_id: str | SequenceNotStr[str], ignore_unknown_id extra_body_fields={"ignoreUnknownIds": ignore_unknown_ids}, headers=self._beta_version_header(), ) + + async def startup(self, request: StartupRequest) -> CheckinResponse: + """`Report extractor startup `_ + + Reports that the extractor has (re)started, along with its current task configuration. + This closes any currently running tasks with an error. + + Note: + This is normally only called by extractor implementations as part of the + integrations startup protocol, not by typical SDK consumers. + + Args: + request (StartupRequest): The startup event to report. + + Returns: + CheckinResponse: The integration's latest config revision. + + Examples: + + Report extractor startup: + + >>> from cognite.client import CogniteClient + >>> from cognite.client.data_classes.integrations import Extractor, StartupRequest + >>> client = CogniteClient() + >>> req = StartupRequest( + ... external_id="my-integration", + ... extractor=Extractor(external_id="cognite-simple-influxdb-extractor"), + ... ) + >>> res = client.integrations.startup(req) + """ + self._warning.warn() + response = await self._post( + f"{self._RESOURCE_PATH}/startup", + json=request.dump(camel_case=True), + headers=self._beta_version_header(), + semaphore=self._get_semaphore("write"), + ) + return CheckinResponse._load(response.json()) + + async def checkin(self, request: CheckinRequest) -> CheckinResponse: + """`Check in with the integrations service `_ + + Called periodically by extractors to signal that they are still alive, and to report task + start/stop events and errors that have occurred since the last check-in. + + Note: + This is normally only called by extractor implementations as part of the + integrations check-in protocol, not by typical SDK consumers. + + Args: + request (CheckinRequest): The check-in event to report. + + Returns: + CheckinResponse: The integration's latest config revision. + + Examples: + + Check in with no updates: + + >>> from cognite.client import CogniteClient + >>> from cognite.client.data_classes.integrations import CheckinRequest + >>> client = CogniteClient() + >>> res = client.integrations.checkin(CheckinRequest(external_id="my-integration")) + """ + self._warning.warn() + response = await self._post( + f"{self._RESOURCE_PATH}/checkin", + json=request.dump(camel_case=True), + headers=self._beta_version_header(), + semaphore=self._get_semaphore("write"), + ) + return CheckinResponse._load(response.json()) diff --git a/cognite/client/_sync_api/integrations/__init__.py b/cognite/client/_sync_api/integrations/__init__.py index 155f1d5f33..2e371c6b30 100644 --- a/cognite/client/_sync_api/integrations/__init__.py +++ b/cognite/client/_sync_api/integrations/__init__.py @@ -1,6 +1,6 @@ """ =============================================================================== -e76b177ed7a693a6551979741ff94979 +b875794d0b27a94c97637678fc1ccd6a This file is auto-generated from the Async API modules, - do not edit manually! =============================================================================== """ @@ -17,6 +17,7 @@ from cognite.client._sync_api.integrations.errors import SyncIntegrationErrorsAPI from cognite.client._sync_api.integrations.tasks import SyncIntegrationTasksAPI from cognite.client._sync_api_client import SyncAPIClient +from cognite.client.data_classes.integrations.checkin import CheckinRequest, CheckinResponse, StartupRequest from cognite.client.data_classes.integrations.integrations import ( Integration, IntegrationList, @@ -207,3 +208,63 @@ def delete(self, external_id: str | SequenceNotStr[str], ignore_unknown_ids: boo return run_sync( self.__async_client.integrations.delete(external_id=external_id, ignore_unknown_ids=ignore_unknown_ids) ) + + def startup(self, request: StartupRequest) -> CheckinResponse: + """ + `Report extractor startup `_ + + Reports that the extractor has (re)started, along with its current task configuration. + This closes any currently running tasks with an error. + + Note: + This is normally only called by extractor implementations as part of the + integrations startup protocol, not by typical SDK consumers. + + Args: + request (StartupRequest): The startup event to report. + + Returns: + CheckinResponse: The integration's latest config revision. + + Examples: + + Report extractor startup: + + >>> from cognite.client import CogniteClient + >>> from cognite.client.data_classes.integrations import Extractor, StartupRequest + >>> client = CogniteClient() + >>> req = StartupRequest( + ... external_id="my-integration", + ... extractor=Extractor(external_id="cognite-simple-influxdb-extractor"), + ... ) + >>> res = client.integrations.startup(req) + """ + return run_sync(self.__async_client.integrations.startup(request=request)) + + def checkin(self, request: CheckinRequest) -> CheckinResponse: + """ + `Check in with the integrations service `_ + + Called periodically by extractors to signal that they are still alive, and to report task + start/stop events and errors that have occurred since the last check-in. + + Note: + This is normally only called by extractor implementations as part of the + integrations check-in protocol, not by typical SDK consumers. + + Args: + request (CheckinRequest): The check-in event to report. + + Returns: + CheckinResponse: The integration's latest config revision. + + Examples: + + Check in with no updates: + + >>> from cognite.client import CogniteClient + >>> from cognite.client.data_classes.integrations import CheckinRequest + >>> client = CogniteClient() + >>> res = client.integrations.checkin(CheckinRequest(external_id="my-integration")) + """ + return run_sync(self.__async_client.integrations.checkin(request=request)) diff --git a/cognite/client/data_classes/integrations/__init__.py b/cognite/client/data_classes/integrations/__init__.py index 569d2c4767..63f866fe9e 100644 --- a/cognite/client/data_classes/integrations/__init__.py +++ b/cognite/client/data_classes/integrations/__init__.py @@ -6,6 +6,13 @@ ActionWrite, ActionWriteList, ) +from cognite.client.data_classes.integrations.checkin import ( + CheckinRequest, + CheckinResponse, + ErrorWithTask, + StartupRequest, + TaskUpdate, +) from cognite.client.data_classes.integrations.config import ( ConfigRevision, ConfigRevisionMetadata, @@ -36,10 +43,13 @@ "ActionList", "ActionWrite", "ActionWriteList", + "CheckinRequest", + "CheckinResponse", "ConfigRevision", "ConfigRevisionMetadata", "ConfigRevisionMetadataList", "ConfigRevisionWrite", + "ErrorWithTask", "Extractor", "Integration", "IntegrationError", @@ -48,8 +58,10 @@ "IntegrationUpdate", "IntegrationWrite", "IntegrationWriteList", + "StartupRequest", "SyncResult", "Task", "TaskHistory", "TaskHistoryList", + "TaskUpdate", ] diff --git a/cognite/client/data_classes/integrations/checkin.py b/cognite/client/data_classes/integrations/checkin.py new file mode 100644 index 0000000000..b6caeda3d5 --- /dev/null +++ b/cognite/client/data_classes/integrations/checkin.py @@ -0,0 +1,194 @@ +from __future__ import annotations + +from typing import Any, Literal, TypeAlias + +from typing_extensions import Self + +from cognite.client.data_classes._base import CogniteResource +from cognite.client.data_classes.integrations.integrations import ActiveConfigRevision, Extractor, Task + +TaskUpdateType: TypeAlias = Literal["started", "ended"] + + +class TaskUpdate(CogniteResource): + """A single start/stop event for a task, reported by the extractor on check-in. + + Args: + type (TaskUpdateType): Whether the task started or ended. + name (str): Name of the task being updated. + timestamp (int): Time of the event, in milliseconds since epoch. + message (str | None): Optional message tied to the task run. + """ + + def __init__(self, type: TaskUpdateType, name: str, timestamp: int, message: str | None = None) -> None: + self.type = type + self.name = name + self.timestamp = timestamp + self.message = message + + @classmethod + def _load(cls, resource: dict[str, Any]) -> Self: + return cls( + type=resource["type"], + name=resource["name"], + timestamp=resource["timestamp"], + message=resource.get("message"), + ) + + +class ErrorWithTask(CogniteResource): + """An error reported by the extractor as part of a check-in. + + Args: + level (Literal['warning', 'error', 'fatal']): Severity of the error. + description (str): Short description of the error. + start_time (int): Time the error started, in milliseconds since epoch. + details (str | None): Full details of the error, e.g. a stack trace. + task (str | None): Name of the task the error occurred in. Not set if the error applies to the extractor generally. + end_time (int | None): Time the error was resolved, in milliseconds since epoch. Not set while unresolved. + active_config_revision (ActiveConfigRevision | None): The config revision (or "local") active when the error occurred. + """ + + def __init__( + self, + level: Literal["warning", "error", "fatal"], + description: str, + start_time: int, + details: str | None = None, + task: str | None = None, + end_time: int | None = None, + active_config_revision: ActiveConfigRevision | None = None, + ) -> None: + self.level = level + self.description = description + self.start_time = start_time + self.details = details + self.task = task + self.end_time = end_time + self.active_config_revision = active_config_revision + + @classmethod + def _load(cls, resource: dict[str, Any]) -> Self: + return cls( + level=resource["level"], + description=resource["description"], + start_time=resource["startTime"], + details=resource.get("details"), + task=resource.get("task"), + end_time=resource.get("endTime"), + active_config_revision=resource.get("activeConfigRevision"), + ) + + +class StartupRequest(CogniteResource): + """Reported by an extractor on startup: general information about the extractor, and an + indication that it has (re)started. This closes any currently running tasks with an error. + + Note: + This is normally only sent by extractor implementations as part of the integrations + startup protocol, not by typical SDK consumers. + + Args: + external_id (str): External id of the integration. + extractor (Extractor): The extractor reporting the startup event. + tasks (list[Task] | None): The tasks configured for this extractor. + active_config_revision (ActiveConfigRevision | None): The config revision (or "local") currently active. + timestamp (int | None): Time of the startup event, in milliseconds since epoch. + """ + + def __init__( + self, + external_id: str, + extractor: Extractor, + tasks: list[Task] | None = None, + active_config_revision: ActiveConfigRevision | None = None, + timestamp: int | None = None, + ) -> None: + self.external_id = external_id + self.extractor = extractor + self.tasks = tasks + self.active_config_revision = active_config_revision + self.timestamp = timestamp + + def dump(self, camel_case: bool = True) -> dict[str, Any]: + result = super().dump(camel_case) + result["extractor"] = self.extractor.dump(camel_case) + if self.tasks is not None: + result["tasks"] = [task.dump(camel_case) for task in self.tasks] + return result + + @classmethod + def _load(cls, resource: dict[str, Any]) -> Self: + return cls( + external_id=resource["externalId"], + extractor=Extractor._load(resource["extractor"]), + tasks=[Task._load(task) for task in resource["tasks"]] if resource.get("tasks") is not None else None, + active_config_revision=resource.get("activeConfigRevision"), + timestamp=resource.get("timestamp"), + ) + + +class CheckinRequest(CogniteResource): + """Reported periodically by an extractor to signal it is still alive, and to report task + start/stop events and errors that have occurred since the last check-in. + + Note: + This is normally only sent by extractor implementations as part of the integrations + check-in protocol, not by typical SDK consumers. + + Args: + external_id (str): External id of the integration. + task_events (list[TaskUpdate] | None): Task start/stop events since the last check-in. + errors (list[ErrorWithTask] | None): Errors reported since the last check-in. + """ + + def __init__( + self, + external_id: str, + task_events: list[TaskUpdate] | None = None, + errors: list[ErrorWithTask] | None = None, + ) -> None: + self.external_id = external_id + self.task_events = task_events + self.errors = errors + + def dump(self, camel_case: bool = True) -> dict[str, Any]: + result = super().dump(camel_case) + if self.task_events is not None: + key = "taskEvents" if camel_case else "task_events" + result[key] = [event.dump(camel_case) for event in self.task_events] + if self.errors is not None: + result["errors"] = [error.dump(camel_case) for error in self.errors] + return result + + @classmethod + def _load(cls, resource: dict[str, Any]) -> Self: + return cls( + external_id=resource["externalId"], + task_events=[TaskUpdate._load(event) for event in resource["taskEvents"]] + if resource.get("taskEvents") is not None + else None, + errors=[ErrorWithTask._load(error) for error in resource["errors"]] + if resource.get("errors") is not None + else None, + ) + + +class CheckinResponse(CogniteResource): + """Response returned from both startup and check-in, containing the latest config revision. + + Args: + external_id (str): External id of the integration. + last_config_revision (int | None): The latest stored configuration revision for this integration. + """ + + def __init__(self, external_id: str, last_config_revision: int | None = None) -> None: + self.external_id = external_id + self.last_config_revision = last_config_revision + + @classmethod + def _load(cls, resource: dict[str, Any]) -> Self: + return cls( + external_id=resource["externalId"], + last_config_revision=resource.get("lastConfigRevision"), + ) diff --git a/cognite/client/utils/_url.py b/cognite/client/utils/_url.py index ffff73943c..3dfbf7da05 100644 --- a/cognite/client/utils/_url.py +++ b/cognite/client/utils/_url.py @@ -60,6 +60,8 @@ "annotations/suggest", "extpipes/config/revert", "integrations/actions/cancel", + "integrations/checkin", + "integrations/startup", "transformations/cancel", "transformations/notifications", "transformations/run", diff --git a/tests/tests_unit/test_api/test_integrations/test_integrations.py b/tests/tests_unit/test_api/test_integrations/test_integrations.py index 170aae84e7..aac668efd7 100644 --- a/tests/tests_unit/test_api/test_integrations/test_integrations.py +++ b/tests/tests_unit/test_api/test_integrations/test_integrations.py @@ -6,11 +6,16 @@ from cognite.client import AsyncCogniteClient, CogniteClient from cognite.client.data_classes.integrations import ( + CheckinRequest, + CheckinResponse, + ErrorWithTask, Extractor, Integration, IntegrationList, IntegrationUpdate, IntegrationWrite, + StartupRequest, + TaskUpdate, ) from tests.utils import get_url, jsgz_load @@ -87,6 +92,56 @@ def test_retrieve( body = jsgz_load(httpx_mock.get_requests()[0].content) assert body == {"items": [{"externalId": "my-integration"}], "ignoreUnknownIds": False} + def test_startup( + self, cognite_client: CogniteClient, async_client: AsyncCogniteClient, httpx_mock: HTTPXMock + ) -> None: + httpx_mock.add_response( + method="POST", + url=get_url(async_client.integrations, "/integrations/startup"), + json={"externalId": "my-integration", "lastConfigRevision": 3}, + ) + request = StartupRequest( + external_id="my-integration", + extractor=Extractor(external_id="cognite-simple-influxdb-extractor", version="1.0.0"), + ) + + res = cognite_client.integrations.startup(request) + + assert isinstance(res, CheckinResponse) + assert res.last_config_revision == 3 + + body = jsgz_load(httpx_mock.get_requests()[0].content) + assert body == { + "externalId": "my-integration", + "extractor": {"externalId": "cognite-simple-influxdb-extractor", "version": "1.0.0"}, + } + + def test_checkin( + self, cognite_client: CogniteClient, async_client: AsyncCogniteClient, httpx_mock: HTTPXMock + ) -> None: + httpx_mock.add_response( + method="POST", + url=get_url(async_client.integrations, "/integrations/checkin"), + json={"externalId": "my-integration", "lastConfigRevision": 4}, + ) + request = CheckinRequest( + external_id="my-integration", + task_events=[TaskUpdate(type="started", name="poll", timestamp=100)], + errors=[ErrorWithTask(level="warning", description="Slow response", start_time=100)], + ) + + res = cognite_client.integrations.checkin(request) + + assert isinstance(res, CheckinResponse) + assert res.last_config_revision == 4 + + body = jsgz_load(httpx_mock.get_requests()[0].content) + assert body == { + "externalId": "my-integration", + "taskEvents": [{"type": "started", "name": "poll", "timestamp": 100}], + "errors": [{"level": "warning", "description": "Slow response", "startTime": 100}], + } + def test_update( self, cognite_client: CogniteClient, async_client: AsyncCogniteClient, httpx_mock: HTTPXMock ) -> None: diff --git a/tests/tests_unit/test_api_client.py b/tests/tests_unit/test_api_client.py index 84faaa9b64..37282bb810 100644 --- a/tests/tests_unit/test_api_client.py +++ b/tests/tests_unit/test_api_client.py @@ -1790,6 +1790,8 @@ async def test_is_retryable_resource_api_endpoints(self, method: str, path: str, ("POST", "https://api.cognitedata.com/api/v1/projects/bla/integrations/actions", False), ("POST", "https://api.cognitedata.com/api/v1/projects/bla/integrations/actions/byids", True), ("POST", "https://api.cognitedata.com/api/v1/projects/bla/integrations/actions/cancel", False), + ("POST", "https://api.cognitedata.com/api/v1/projects/bla/integrations/startup", False), + ("POST", "https://api.cognitedata.com/api/v1/projects/bla/integrations/checkin", False), # Transformations ("POST", "https://api.cognitedata.com/api/v1/projects/bla/transformations", False), ("POST", "https://api.cognitedata.com/api/v1/projects/bla/transformations/filter", True), diff --git a/tests/tests_unit/test_data_classes/test_integrations.py b/tests/tests_unit/test_data_classes/test_integrations.py index 9c1a4af29f..c464af3bc5 100644 --- a/tests/tests_unit/test_data_classes/test_integrations.py +++ b/tests/tests_unit/test_data_classes/test_integrations.py @@ -2,12 +2,17 @@ from cognite.client.data_classes.integrations import ( Action, + CheckinRequest, + CheckinResponse, ConfigRevision, + ErrorWithTask, Extractor, Integration, IntegrationError, IntegrationUpdate, + StartupRequest, Task, + TaskUpdate, ) from cognite.client.data_classes.integrations.tasks import SyncResult, TaskHistory @@ -204,3 +209,55 @@ def test_task_load_dump() -> None: "targets": ["cdf://cluster/project/timeseries/my-ts"], } assert Task._load(dumped).dump(camel_case=True) == dumped + + +def test_task_update_load_dump() -> None: + dumped = {"type": "started", "name": "poll", "timestamp": 100, "message": "Task started"} + assert TaskUpdate._load(dumped).dump(camel_case=True) == dumped + + +def test_error_with_task_load_dump() -> None: + dumped = { + "level": "fatal", + "description": "Something went very wrong", + "startTime": 100, + "details": "Traceback ...", + "task": "poll", + "endTime": 200, + "activeConfigRevision": 2, + } + assert ErrorWithTask._load(dumped).dump(camel_case=True) == dumped + + +def test_startup_request_load_dump() -> None: + dumped = { + "externalId": "my-integration", + "extractor": {"externalId": "cognite-simple-influxdb-extractor", "version": "1.0.0"}, + "tasks": [{"type": "continuous", "name": "poll", "action": False}], + "activeConfigRevision": "local", + "timestamp": 100, + } + loaded = StartupRequest._load(dumped) + + assert loaded.tasks is not None + assert loaded.tasks[0].name == "poll" + assert loaded.dump(camel_case=True) == dumped + + +def test_checkin_request_load_dump() -> None: + dumped = { + "externalId": "my-integration", + "taskEvents": [{"type": "ended", "name": "poll", "timestamp": 200}], + "errors": [{"level": "error", "description": "Oops", "startTime": 100}], + } + loaded = CheckinRequest._load(dumped) + + assert loaded.task_events is not None + assert loaded.errors is not None + assert loaded.task_events[0].type == "ended" + assert loaded.dump(camel_case=True) == dumped + + +def test_checkin_response_load_dump() -> None: + dumped = {"externalId": "my-integration", "lastConfigRevision": 5} + assert CheckinResponse._load(dumped).dump(camel_case=True) == dumped