cognite-toolkit 0.7.43__py3-none-any.whl → 0.7.44__py3-none-any.whl

cognite_toolkit/_cdf_tk/client/data_classes/annotation.py

@@ -0,0 +1,79 @@
+from typing import Any, Literal, TypeAlias
+
+from pydantic import Field, JsonValue
+
+from cognite_toolkit._cdf_tk.client.data_classes.base import BaseModelObject, RequestUpdateable, ResponseResource
+
+from .identifiers import InternalId
+
+AnnotationStatus: TypeAlias = Literal["suggested", "rejected", "approved"]
+AnnotationType: TypeAlias = Literal[
+    "diagrams.AssetLink",
+    "diagrams.FileLink",
+    "diagrams.InstanceLink",
+    "diagrams.Junction",
+    "diagrams.Line",
+    "diagrams.UnhandledSymbolObject",
+    "diagrams.UnhandledTextObject",
+    "documents.ExtractedText",
+    "forms.Detection",
+    "images.AssetLink",
+    "images.Classification",
+    "images.InstanceLink",
+    "images.KeypointCollection",
+    "images.ObjectDetection",
+    "images.TextRegion",
+    "isoplan.IsoPlanAnnotation",
+    "pointcloud.BoundingVolume",
+]
+
+
+class Annotation(BaseModelObject):
+    annotated_resource_type: str
+    annotated_resource_id: int
+    annotation_type: AnnotationType
+    creating_app: str
+    creating_app_version: str
+    creating_user: str | None
+    data: dict[str, JsonValue]
+    status: AnnotationStatus
+
+
+class AnnotationRequest(Annotation, RequestUpdateable):
+    """Request data class for annotations."""
+
+    # The 'id' field is not part of the request when creating a new resource,
+    # but is needed when updating an existing resource.
+    id: int | None = Field(default=None, exclude=True)
+
+    def as_id(self) -> InternalId:
+        if self.id is None:
+            raise ValueError("Cannot convert AnnotationRequest to InternalId when id is None")
+        return InternalId(id=self.id)
+
+    def as_update(self, mode: Literal["patch", "replace"]) -> dict[str, Any]:
+        """Converts the request to an update payload for the API."""
+        if self.id is None:
+            raise ValueError("id must be provided to create an update dictionary")
+        return {
+            "id": self.id,
+            "update": {
+                "annotationType": {"set": self.annotation_type},
+                "data": {"set": self.data},
+                "status": {"set": self.status},
+            },
+        }
+
+
+class AnnotationResponse(Annotation, ResponseResource[AnnotationRequest]):
+    """Response data class for annotations."""
+
+    id: int
+    created_time: int
+    last_updated_time: int
+
+    def as_id(self) -> InternalId:
+        return InternalId(id=self.id)
+
+    def as_request_resource(self) -> AnnotationRequest:
+        return AnnotationRequest.model_validate(self.dump(), extra="ignore")

cognite_toolkit/_cdf_tk/client/data_classes/base.py

@@ -20,11 +20,21 @@ class BaseModelObject(BaseModel):
     # We allow extra fields to support forward compatibility.
     model_config = ConfigDict(alias_generator=to_camel, extra="allow", populate_by_name=True)
 
-    def dump(self, camel_case: bool = True) -> dict[str, Any]:
+    def dump(self, camel_case: bool = True, exclude_extra: bool = False) -> dict[str, Any]:
         """Dump the resource to a dictionary.
 
-        This is the default serialization method for request resources.
+        Args:
+            camel_case (bool): Whether to use camelCase for the keys. Default is True.
+            exclude_extra (bool): Whether to exclude extra fields not defined in the model. Default is False.
+
         """
+        if exclude_extra:
+            return self.model_dump(
+                mode="json",
+                by_alias=camel_case,
+                exclude_unset=True,
+                exclude=set(self.__pydantic_extra__) if self.__pydantic_extra__ else None,
+            )
         return self.model_dump(mode="json", by_alias=camel_case, exclude_unset=True)
 
     @classmethod
@@ -39,7 +49,7 @@ class Identifier(BaseModelObject):
 
     model_config = ConfigDict(alias_generator=to_camel, extra="ignore", populate_by_name=True, frozen=True)
 
-    def dump(self, camel_case: bool = True) -> dict[str, Any]:
+    def dump(self, camel_case: bool = True, exclude_extra: bool = False) -> dict[str, Any]:
         """Dump the resource to a dictionary.
 
         This is the default serialization method for request resources.

cognite_toolkit/_cdf_tk/client/data_classes/data_modeling/__init__.py

@@ -42,6 +42,15 @@ from ._indexes import (
     IndexDefinition,
     InvertedIndex,
 )
+from ._instance import (
+    EdgeRequest,
+    EdgeResponse,
+    InstanceDefinition,
+    InstanceResponseDefinition,
+    InstanceSource,
+    NodeRequest,
+    NodeResponse,
+)
 from ._references import (
     ContainerConstraintReference,
     ContainerDirectReference,
@@ -99,6 +108,8 @@ __all__ = [
     "DateProperty",
     "DirectNodeRelation",
     "EdgeProperty",
+    "EdgeRequest",
+    "EdgeResponse",
     "EnumProperty",
     "EnumValue",
     "FileCDFExternalIdReference",
@@ -108,6 +119,9 @@ __all__ = [
     "Index",
     "IndexAdapter",
     "IndexDefinition",
+    "InstanceDefinition",
+    "InstanceResponseDefinition",
+    "InstanceSource",
     "Int32Property",
     "Int64Property",
     "InvertedIndex",
@@ -117,6 +131,8 @@ __all__ = [
     "MultiReverseDirectRelationPropertyRequest",
     "MultiReverseDirectRelationPropertyResponse",
     "NodeReference",
+    "NodeRequest",
+    "NodeResponse",
     "PropertyTypeDefinition",
     "RequiresConstraintDefinition",
     "ReverseDirectRelationProperty",

cognite_toolkit/_cdf_tk/client/data_classes/data_modeling/_instance.py

@@ -0,0 +1,143 @@
+from abc import ABC
+from typing import Any, Generic, Literal
+
+from pydantic import JsonValue, field_serializer, field_validator
+
+from cognite_toolkit._cdf_tk.client.data_classes.base import (
+    BaseModelObject,
+    RequestResource,
+    ResponseResource,
+    T_RequestResource,
+)
+
+from ._references import ContainerReference, EdgeReference, NodeReference, ViewReference
+
+
+class InstanceDefinition(BaseModelObject, ABC):
+    """Base class for node and edge instances."""
+
+    instance_type: str  # "node" | "edge"
+    space: str
+    external_id: str
+
+
+class InstanceSource(BaseModelObject):
+    source: ViewReference | ContainerReference
+    properties: dict[str, JsonValue] | None = None
+
+
+class InstanceRequestDefinition(InstanceDefinition, RequestResource, ABC):
+    existing_version: int | None = None
+    sources: list[InstanceSource]
+
+
+class InstanceResponseDefinition(InstanceDefinition, ResponseResource, Generic[T_RequestResource], ABC):
+    version: int
+    created_time: int
+    last_updated_time: int
+    deleted_time: int | None = None
+    properties: dict[ViewReference | ContainerReference, dict[str, JsonValue]] | None = None
+
+    @field_validator("properties", mode="before")
+    def parse_reference(cls, value: Any) -> Any:
+        if not isinstance(value, dict):
+            return value
+        parsed: dict[ViewReference | ContainerReference, dict[str, Any]] = {}
+        for space, inner_dict in value.items():
+            if not isinstance(inner_dict, dict) or not isinstance(space, str):
+                raise ValueError(
+                    f"Invalid properties format expected dict[str, dict[...]], got: dict[{type(space).__name__}, {type(inner_dict).__name__}]"
+                )
+            for view_or_container_identifier, prop in inner_dict.items():
+                if not isinstance(view_or_container_identifier, str):
+                    raise ValueError(
+                        "Invalid properties format expected dict[str, dict[str, ...]]], "
+                        f"got: dict[{type(space).__name__}, "
+                        f"dict[{type(view_or_container_identifier).__name__}, ...]]"
+                    )
+                source_ref: ViewReference | ContainerReference
+                if "/" in view_or_container_identifier:
+                    external_id, version = view_or_container_identifier.split("/", 1)
+                    source_ref = ViewReference(space=space, external_id=external_id, version=version)
+                else:
+                    source_ref = ContainerReference(space=space, external_id=view_or_container_identifier)
+                parsed[source_ref] = prop
+        return parsed
+
+    @field_serializer("properties", mode="plain")
+    def serialize_properties(self, value: dict[ViewReference | ContainerReference, dict[str, Any]] | None) -> Any:
+        if value is None:
+            return None
+        serialized: dict[str, dict[str, Any]] = {}
+        for source_ref, props in value.items():
+            space = source_ref.space
+            if space not in serialized:
+                serialized[space] = {}
+            if isinstance(source_ref, ViewReference):
+                identifier = f"{source_ref.external_id}/{source_ref.version}"
+            else:
+                identifier = source_ref.external_id
+            serialized[space][identifier] = props
+        return serialized
+
+
+class NodeRequest(InstanceRequestDefinition):
+    """A node request resource."""
+
+    instance_type: Literal["node"] = "node"
+    type: NodeReference | None = None
+
+    def as_id(self) -> NodeReference:
+        return NodeReference(space=self.space, external_id=self.external_id)
+
+
+class EdgeRequest(InstanceRequestDefinition):
+    """An edge request resource."""
+
+    instance_type: Literal["edge"] = "edge"
+    type: NodeReference
+    start_node: NodeReference
+    end_node: NodeReference
+
+    def as_id(self) -> EdgeReference:
+        return EdgeReference(space=self.space, external_id=self.external_id)
+
+
+class NodeResponse(InstanceResponseDefinition[NodeRequest]):
+    """A node response from the API."""
+
+    instance_type: Literal["node"] = "node"
+    type: NodeReference | None = None
+
+    def as_id(self) -> NodeReference:
+        return NodeReference(space=self.space, external_id=self.external_id)
+
+    def as_request_resource(self) -> NodeRequest:
+        dumped = self.dump()
+        if self.properties:
+            dumped["sources"] = [
+                InstanceSource(source=source_ref, properties=props) for source_ref, props in self.properties.items()
+            ]
+        dumped["existingVersion"] = dumped.pop("version", None)
+        return NodeRequest.model_validate(dumped, extra="ignore")
+
+
+class EdgeResponse(InstanceResponseDefinition[EdgeRequest]):
+    """An edge response from the API."""
+
+    instance_type: Literal["edge"] = "edge"
+    type: NodeReference
+    start_node: NodeReference
+    end_node: NodeReference
+
+    def as_id(self) -> EdgeReference:
+        return EdgeReference(space=self.space, external_id=self.external_id)
+
+    def as_request_resource(self) -> EdgeRequest:
+        dumped = self.dump()
+        if self.properties:
+            dumped["sources"] = [
+                InstanceSource(source=source_ref, properties=props) for source_ref, props in self.properties.items()
+            ]
+        dumped["existingVersion"] = dumped.pop("version", None)
+        return EdgeRequest.model_validate(dumped, extra="ignore")

cognite_toolkit/_cdf_tk/client/data_classes/data_modeling/_references.py

@@ -48,6 +48,14 @@ class NodeReference(Identifier):
         return f"{self.space}:{self.external_id}"
 
 
+class EdgeReference(Identifier):
+    space: str
+    external_id: str
+
+    def __str__(self) -> str:
+        return f"{self.space}:{self.external_id}"
+
+
 class ContainerDirectReference(Identifier):
     source: ContainerReference
     identifier: str

cognite_toolkit/_cdf_tk/client/data_classes/dataset.py

@@ -0,0 +1,35 @@
+from typing import ClassVar
+
+from cognite_toolkit._cdf_tk.client.data_classes.base import (
+    BaseModelObject,
+    RequestUpdateable,
+    ResponseResource,
+)
+
+from .identifiers import ExternalId
+
+
+class DataSet(BaseModelObject):
+    external_id: str | None = None
+    name: str | None = None
+    description: str | None = None
+    metadata: dict[str, str] | None = None
+    write_protected: bool = False
+
+    def as_id(self) -> ExternalId:
+        if self.external_id is None:
+            raise ValueError("Cannot convert DataSet to ExternalId when external_id is None")
+        return ExternalId(external_id=self.external_id)
+
+
+class DataSetRequest(DataSet, RequestUpdateable):
+    container_fields: ClassVar[frozenset[str]] = frozenset({"metadata"})
+
+
+class DataSetResponse(DataSet, ResponseResource[DataSetRequest]):
+    id: int
+    created_time: int
+    last_updated_time: int
+
+    def as_request_resource(self) -> DataSetRequest:
+        return DataSetRequest.model_validate(self.dump(), extra="ignore")

cognite_toolkit/_cdf_tk/client/data_classes/extraction_pipeline.py

@@ -0,0 +1,59 @@
+from typing import ClassVar
+
+from cognite_toolkit._cdf_tk.client.data_classes.base import (
+    BaseModelObject,
+    RequestUpdateable,
+    ResponseResource,
+)
+
+from .identifiers import ExternalId
+
+
+class RawTable(BaseModelObject):
+    db_name: str
+    table_name: str
+
+
+class Contact(BaseModelObject):
+    name: str | None = None
+    email: str | None = None
+    role: str | None = None
+    send_notification: bool | None = None
+
+
+class NotificationConfig(BaseModelObject):
+    allowed_not_seen_range_in_minutes: int | None = None
+
+
+class ExtractionPipeline(BaseModelObject):
+    external_id: str
+    name: str
+    description: str | None = None
+    data_set_id: int
+    raw_tables: list[RawTable] | None = None
+    schedule: str | None = None
+    contacts: list[Contact] | None = None
+    metadata: dict[str, str] | None = None
+    source: str | None = None
+    documentation: str | None = None
+    notification_config: NotificationConfig | None = None
+    created_by: str | None = None
+
+    def as_id(self) -> ExternalId:
+        return ExternalId(external_id=self.external_id)
+
+
+class ExtractionPipelineRequest(ExtractionPipeline, RequestUpdateable):
+    container_fields: ClassVar[frozenset[str]] = frozenset({"raw_tables", "contacts", "metadata"})
+    non_nullable_fields: ClassVar[frozenset[str]] = frozenset(
+        {"documentation", "source", "notification_config", "schedule", "description"}
+    )
+
+
+class ExtractionPipelineResponse(ExtractionPipeline, ResponseResource[ExtractionPipelineRequest]):
+    id: int
+    created_time: int
+    last_updated_time: int
+
+    def as_request_resource(self) -> ExtractionPipelineRequest:
+        return ExtractionPipelineRequest.model_validate(self.dump(), extra="ignore")

cognite_toolkit/_cdf_tk/client/data_classes/filemetadata.py

@@ -1,6 +1,6 @@
 from typing import ClassVar, Literal
 
-from pydantic import JsonValue
+from pydantic import Field, JsonValue
 
 from cognite_toolkit._cdf_tk.client.data_classes.base import BaseModelObject, RequestUpdateable, ResponseResource
 
@@ -32,6 +32,10 @@ class FileMetadata(BaseModelObject):
 class FileMetadataRequest(FileMetadata, RequestUpdateable):
     container_fields: ClassVar[frozenset[str]] = frozenset({"metadata", "labels", "asset_ids", "security_categories"})
     non_nullable_fields: ClassVar[frozenset[str]] = frozenset({"asset_ids", "security_categories"})
+    # This field is not part of the request when creating or updating a resource
+    # but we added it here for convenience so that it is available when converting
+    # from response to request.
+    instance_id: NodeReference | None = Field(default=None, exclude=True)
 
 
 class FileMetadataResponse(FileMetadata, ResponseResource[FileMetadataRequest]):
@@ -41,6 +45,8 @@ class FileMetadataResponse(FileMetadata, ResponseResource[FileMetadataRequest]):
     uploaded: bool
     id: int
     instance_id: NodeReference | None = None
+    # This field is required in the upload endpoint response, but not in any other file metadata response
+    upload_url: str | None = None
 
     def as_request_resource(self) -> FileMetadataRequest:
         return FileMetadataRequest.model_validate(self.dump(), extra="ignore")

cognite_toolkit/_cdf_tk/client/data_classes/hosted_extractor_destination.py

@@ -0,0 +1,34 @@
+from cognite_toolkit._cdf_tk.client.data_classes.base import (
+    BaseModelObject,
+    RequestUpdateable,
+    ResponseResource,
+)
+
+from .identifiers import ExternalId
+
+
+class Credentials(BaseModelObject):
+    nonce: str
+
+
+class HostedExtractorDestination(BaseModelObject):
+    external_id: str
+    target_data_set_id: int | None = None
+
+    def as_id(self) -> ExternalId:
+        return ExternalId(external_id=self.external_id)
+
+
+class HostedExtractorDestinationRequest(HostedExtractorDestination, RequestUpdateable):
+    credentials: Credentials | None = None
+
+
+class HostedExtractorDestinationResponse(
+    HostedExtractorDestination, ResponseResource[HostedExtractorDestinationRequest]
+):
+    session_id: int | None = None
+    created_time: int
+    last_updated_time: int
+
+    def as_request_resource(self) -> HostedExtractorDestinationRequest:
+        return HostedExtractorDestinationRequest.model_validate(self.dump(), extra="ignore")

cognite_toolkit/_cdf_tk/client/data_classes/hosted_extractor_job.py

@@ -0,0 +1,134 @@
+from typing import Annotated, Literal
+
+from pydantic import Field, JsonValue
+
+from cognite_toolkit._cdf_tk.client.data_classes.base import (
+    BaseModelObject,
+    RequestUpdateable,
+    ResponseResource,
+)
+
+from .identifiers import ExternalId
+
+
+class JobFormatDefinition(BaseModelObject):
+    type: str
+
+
+class PrefixConfig(BaseModelObject):
+    from_topic: bool | None = None
+    prefix: str | None = None
+
+
+class SpaceRef(BaseModelObject):
+    space: str
+
+
+class CogniteFormat(JobFormatDefinition):
+    type: Literal["cognite"] = "cognite"
+    encoding: str | None = None
+    compression: str | None = None
+    prefix: PrefixConfig | None = None
+    data_models: list[SpaceRef] | None = None
+
+
+class CustomFormat(JobFormatDefinition):
+    type: Literal["custom"] = "custom"
+    encoding: str | None = None
+    compression: str | None = None
+    mapping_id: str
+
+
+class RockwellFormat(JobFormatDefinition):
+    type: Literal["rockwell"] = "rockwell"
+    encoding: str | None = None
+    compression: str | None = None
+    prefix: PrefixConfig | None = None
+    data_models: list[SpaceRef] | None = None
+
+
+class ValueFormat(JobFormatDefinition):
+    type: Literal["value"] = "value"
+    encoding: str | None = None
+    compression: str | None = None
+    prefix: PrefixConfig | None = None
+    data_models: list[SpaceRef] | None = None
+
+
+JobFormat = Annotated[
+    CogniteFormat | CustomFormat | RockwellFormat | ValueFormat,
+    Field(discriminator="type"),
+]
+
+
+class IncrementalLoadDefinition(BaseModelObject):
+    type: str
+
+
+class BodyIncrementalLoad(IncrementalLoadDefinition):
+    type: Literal["body"] = "body"
+    value: str
+
+
+class NextUrlIncrementalLoad(IncrementalLoadDefinition):
+    type: Literal["nextUrl"] = "nextUrl"
+    value: str
+
+
+class HeaderIncrementalLoad(IncrementalLoadDefinition):
+    type: Literal["headerValue"] = "headerValue"
+    key: str
+    value: str
+
+
+class QueryParamIncrementalLoad(IncrementalLoadDefinition):
+    type: Literal["queryParam"] = "queryParam"
+    key: str
+    value: str
+
+
+class MQTTConfig(BaseModelObject):
+    topic_filter: str
+
+
+class KafkaConfig(BaseModelObject):
+    topic: str
+    partitions: int | None = None
+
+
+class RestConfig(BaseModelObject):
+    interval: Literal["5m", "15m", "1h", "6h", "12h", "1d"]
+    path: str
+    method: Literal["get", "post"] | None = None
+    body: JsonValue | None = None
+    query: dict[str, str] | None = None
+    headers: dict[str, str] | None = None
+    incremental_load: BodyIncrementalLoad | HeaderIncrementalLoad | QueryParamIncrementalLoad | None = Field(
+        None, discriminator="type"
+    )
+    pagination: (
+        BodyIncrementalLoad | NextUrlIncrementalLoad | HeaderIncrementalLoad | QueryParamIncrementalLoad | None
+    ) = Field(None, discriminator="type")
+
+
+class HostedExtractorJob(BaseModelObject):
+    external_id: str
+    destination_id: str
+    source_id: str
+    format: JobFormat
+    config: MQTTConfig | KafkaConfig | RestConfig
+
+    def as_id(self) -> ExternalId:
+        return ExternalId(external_id=self.external_id)
+
+
+class HostedExtractorJobRequest(HostedExtractorJob, RequestUpdateable): ...
+
+
+class HostedExtractorJobResponse(HostedExtractorJob, ResponseResource[HostedExtractorJobRequest]):
+    status: str | None = None
+    created_time: int
+    last_updated_time: int
+
+    def as_request_resource(self) -> HostedExtractorJobRequest:
+        return HostedExtractorJobRequest.model_validate(self.dump(), extra="ignore")

cognite_toolkit/_cdf_tk/client/data_classes/hosted_extractor_mapping.py

@@ -0,0 +1,72 @@
+from typing import Annotated, Literal
+
+from pydantic import Field
+
+from cognite_toolkit._cdf_tk.client.data_classes.base import (
+    BaseModelObject,
+    RequestUpdateable,
+    ResponseResource,
+)
+
+from .identifiers import ExternalId
+
+
+class Mapping(BaseModelObject):
+    expression: str
+
+
+class MappingInputDefinition(BaseModelObject):
+    type: str
+
+
+class ProtobufFile(BaseModelObject):
+    file_name: str
+    content: str
+
+
+class ProtobufInput(BaseModelObject):
+    type: Literal["protobuf"] = "protobuf"
+    message_name: str
+    files: list[ProtobufFile]
+
+
+class CSVInput(BaseModelObject):
+    type: Literal["csv"] = "csv"
+    delimiter: str = ","
+    custom_keys: list[str] | None = None
+
+
+class XMLInput(BaseModelObject):
+    type: Literal["xml"] = "xml"
+
+
+class JSONInput(BaseModelObject):
+    type: Literal["json"] = "json"
+
+
+MappingInput = Annotated[
+    ProtobufInput | CSVInput | XMLInput | JSONInput,
+    Field(discriminator="type"),
+]
+
+
+class HostedExtractorMapping(BaseModelObject):
+    external_id: str
+    mapping: Mapping
+    published: bool
+
+    def as_id(self) -> ExternalId:
+        return ExternalId(external_id=self.external_id)
+
+
+class HostedExtractorMappingRequest(HostedExtractorMapping, RequestUpdateable):
+    input: MappingInput | None = None
+
+
+class HostedExtractorMappingResponse(HostedExtractorMapping, ResponseResource[HostedExtractorMappingRequest]):
+    input: MappingInput
+    created_time: int
+    last_updated_time: int
+
+    def as_request_resource(self) -> HostedExtractorMappingRequest:
+        return HostedExtractorMappingRequest.model_validate(self.dump(), extra="ignore")