cognite-neat 0.97.3__py3-none-any.whl → 0.99.0__py3-none-any.whl

cognite/neat/{_utils/cdf/data_classes.py → _client/data_classes/data_modeling.py}

@@ -1,26 +1,15 @@
 from abc import ABC, abstractmethod
-from collections.abc import (
-    Hashable,
-    ItemsView,
-    Iterable,
-    Iterator,
-    KeysView,
-    Mapping,
-    MutableMapping,
-    ValuesView,
+from collections.abc import Hashable, ItemsView, Iterable, Iterator, KeysView, Mapping, MutableMapping, ValuesView
+from typing import (
+    Any,
+    TypeVar,
+    cast,
+    final,
 )
-from dataclasses import dataclass
-from typing import Any, TypeVar, cast, final
 
 import pandas as pd
 import yaml
-from cognite.client import CogniteClient
-from cognite.client.data_classes._base import (
-    CogniteResourceList,
-    T_CogniteResource,
-    WriteableCogniteResource,
-    WriteableCogniteResourceList,
-)
+from cognite.client.data_classes._base import T_CogniteResource
 from cognite.client.data_classes.data_modeling import (
     ContainerApply,
     ContainerId,
@@ -33,127 +22,11 @@ from cognite.client.data_classes.data_modeling import (
     ViewId,
 )
 from cognite.client.utils._auxiliary import load_yaml_or_json
-from cognite.client.utils._pandas_helpers import (
-    convert_nullable_int_cols,
-)
-
-# The Table, TableWrite data classes in the Cognite-SDK lacks the database attribute.
-# This is a problem when creating the RawTableLoader that needs the data class to be able to create, update, retrieve
-# and delete tables.
-# This is a reimplemented version of the Table, TableWrite data classes with the database attribute added.
-
-
-@dataclass(frozen=True)
-class RawTableID:
-    table: str
-    database: str
-
-    def as_tuple(self) -> tuple[str, str]:
-        return self.database, self.table
-
-
-class RawTableCore(WriteableCogniteResource["RawTableWrite"], ABC):
-    """A NoSQL database table to store customer data
-
-    Args:
-        name (str | None): Unique name of the table
-    """
-
-    def __init__(
-        self,
-        name: str | None = None,
-        database: str | None = None,
-    ) -> None:
-        self.name = name
-        self.database = database
-
-    def as_id(self) -> RawTableID:
-        if self.name is None or self.database is None:
-            raise ValueError("name and database are required to create a TableID")
-        return RawTableID(table=self.name, database=self.database)
-
-
-class RawTable(RawTableCore):
-    """A NoSQL database table to store customer data.
-    This is the reading version of the Table class, which is used when retrieving a table.
-
-    Args:
-        name (str | None): Unique name of the table
-        created_time (int | None): Time the table was created.
-        cognite_client (CogniteClient | None): The client to associate with this object.
-    """
-
-    def __init__(
-        self,
-        name: str | None = None,
-        database: str | None = None,
-        created_time: int | None = None,
-        cognite_client: CogniteClient | None = None,
-    ) -> None:
-        super().__init__(name, database)
-        self.created_time = created_time
-        self._cognite_client = cast("CogniteClient", cognite_client)
-
-        self._db_name: str | None = None
-
-    def as_write(self) -> "RawTableWrite":
-        """Returns this Table as a TableWrite"""
-        if self.name is None or self.database is None:
-            raise ValueError("name and database are required to create a Table")
-        return RawTableWrite(name=self.name, database=self.database)
-
-
-class RawTableWrite(RawTableCore):
-    """A NoSQL database table to store customer data
-    This is the writing version of the Table class, which is used when creating a table.
-
-    Args:
-        name (str): Unique name of the table
-    """
-
-    def __init__(
-        self,
-        name: str,
-        database: str,
-    ) -> None:
-        super().__init__(name, database)
-
-    @classmethod
-    def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> "RawTableWrite":
-        return cls(resource["name"], resource["database"])
-
-    def as_write(self) -> "RawTableWrite":
-        """Returns this TableWrite instance."""
-        return self
-
-
-class RawTableWriteList(CogniteResourceList[RawTableWrite]):
-    _RESOURCE = RawTableWrite
-
-    def as_ids(self) -> list[RawTableID]:
-        """Returns this TableWriteList as a list of TableIDs"""
-        return [table.as_id() for table in self.data]
-
-
-class RawTableList(
-    WriteableCogniteResourceList[RawTableWrite, RawTable],
-):
-    _RESOURCE = RawTable
-
-    def as_write(self) -> RawTableWriteList:
-        """Returns this TableList as a TableWriteList"""
-        return RawTableWriteList([table.as_write() for table in self.data])
-
-    def as_ids(self) -> list[RawTableID]:
-        """Returns this TableList as a list of TableIDs"""
-        return [table.as_id() for table in self.data]
-
+from cognite.client.utils._pandas_helpers import convert_nullable_int_cols
 
 T_ID = TypeVar("T_ID", bound=Hashable)
 
 
-# Inheriting from dict as we are extending it,
-# ref https://stackoverflow.com/questions/7148419/subclass-dict-userdict-dict-or-abc
 class CogniteResourceDict(dict, MutableMapping[T_ID, T_CogniteResource], ABC):
     """CogniteResource stored in a mapping structure.
 

cognite/neat/{_rules/models/dms/_schema.py → _client/data_classes/schema.py}

@@ -1,4 +1,3 @@
-import json
 import sys
 import warnings
 import zipfile
@@ -6,13 +5,12 @@ from collections import ChainMap, Counter, defaultdict
 from collections.abc import Iterable, MutableMapping
 from dataclasses import Field, dataclass, field, fields
 from pathlib import Path
-from typing import Any, ClassVar, Literal, cast
+from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast
 
 import yaml
 from cognite.client import CogniteClient
 from cognite.client import data_modeling as dm
-from cognite.client.data_classes import DatabaseWrite, DatabaseWriteList, TransformationWrite, TransformationWriteList
-from cognite.client.data_classes.data_modeling import ViewApply
+from cognite.client.data_classes import DatabaseWrite, TransformationWrite
 from cognite.client.data_classes.data_modeling.views import (
     ReverseDirectRelation,
     ReverseDirectRelationApply,
@@ -23,8 +21,14 @@ from cognite.client.data_classes.data_modeling.views import (
     ViewProperty,
     ViewPropertyApply,
 )
-from cognite.client.data_classes.transformations.common import Edges, EdgeType, Nodes, ViewInfo
 
+from cognite.neat._client.data_classes.data_modeling import (
+    CogniteResourceDict,
+    ContainerApplyDict,
+    NodeApplyDict,
+    SpaceApplyDict,
+    ViewApplyDict,
+)
 from cognite.neat._issues import NeatError
 from cognite.neat._issues.errors import (
     NeatYamlError,
@@ -40,17 +44,6 @@ from cognite.neat._issues.warnings import (
     ResourcesDuplicatedWarning,
 )
 from cognite.neat._issues.warnings.user_modeling import DirectRelationMissingSourceWarning
-from cognite.neat._rules.models.data_types import _DATA_TYPE_BY_DMS_TYPE
-from cognite.neat._utils.cdf.data_classes import (
-    CogniteResourceDict,
-    ContainerApplyDict,
-    NodeApplyDict,
-    RawTableWrite,
-    RawTableWriteList,
-    SpaceApplyDict,
-    ViewApplyDict,
-)
-from cognite.neat._utils.cdf.loaders import ViewLoader
 from cognite.neat._utils.rdf_ import get_inheritance_path
 from cognite.neat._utils.text import to_camel
 
@@ -59,6 +52,9 @@ if sys.version_info >= (3, 11):
 else:
     from typing_extensions import Self
 
+if TYPE_CHECKING:
+    from cognite.neat._client import NeatClient
+
 
 @dataclass
 class DMSSchema:
@@ -105,7 +101,7 @@ class DMSSchema:
         return directly_referenced_containers | inherited_referenced_containers
 
     @classmethod
-    def from_model_id(cls, client: CogniteClient, data_model_id: dm.DataModelIdentifier) -> "DMSSchema":
+    def from_model_id(cls, client: "NeatClient", data_model_id: dm.DataModelIdentifier) -> "DMSSchema":
         data_models = client.data_modeling.data_models.retrieve(data_model_id, inline_views=True)
         if len(data_models) == 0:
             raise ValueError(f"Data model {data_model_id} not found")
@@ -115,7 +111,7 @@ class DMSSchema:
     @classmethod
     def from_data_model(
         cls,
-        client: CogniteClient,
+        client: "NeatClient",
         data_model: dm.DataModel[dm.View],
         reference_model: dm.DataModel[dm.View] | None = None,
     ) -> "DMSSchema":
@@ -139,6 +135,8 @@ class DMSSchema:
         Returns:
             DMSSchema: The schema created from the data model.
         """
+        from cognite.neat._client._api.data_modeling_loaders import ViewLoader
+
         views = dm.ViewList(data_model.views)
 
         data_model_write = data_model.as_write()
@@ -157,7 +155,7 @@ class DMSSchema:
             raise ValueError(f"Space(s) {space_read} not found")
         space_write = space_read.as_write()
 
-        view_loader = ViewLoader(client)
+        view_loader = ViewLoader(NeatClient(client))
 
         existing_view_ids = set(views.as_ids())
 
@@ -179,7 +177,9 @@ class DMSSchema:
 
         # We need to include parent views in the schema to make sure that the schema is valid.
         parent_view_ids = {parent for view in views for parent in view.implements or []}
-        parents = view_loader.retrieve_all_parents(list(parent_view_ids - existing_view_ids))
+        parents = view_loader.retrieve(
+            list(parent_view_ids - existing_view_ids), include_ancestor=True, include_connections=True
+        )
         views.extend([parent for parent in parents if parent.as_id() not in existing_view_ids])
 
         # Converting views from read to write format requires to account for parents (implements)
@@ -524,7 +524,7 @@ class DMSSchema:
 
     @classmethod
     def _to_sortable_identifier(cls, item: Any) -> str | tuple[str, str] | tuple[str, str, str]:
-        if isinstance(item, dm.ContainerApply | dm.ViewApply | dm.DataModelApply | dm.NodeApply | RawTableWrite):
+        if isinstance(item, dm.ContainerApply | dm.ViewApply | dm.DataModelApply | dm.NodeApply):
             identifier = item.as_id().as_tuple()
             if len(identifier) == 3 and identifier[2] is None:
                 return identifier[:2]  # type: ignore[misc]
@@ -539,6 +539,10 @@ class DMSSchema:
             raise ValueError(f"Cannot sort item of type {type(item)}")
 
     def validate(self) -> list[NeatError]:
+        # TODO: This type of validation should be done in NeatSession where all the
+        # schema components which are not part of Rules are imported and the model as
+        # the whole is validated.
+
         errors: set[NeatError] = set()
         defined_spaces = self.spaces.copy()
         defined_containers = self.containers.copy()
@@ -708,6 +712,13 @@ class DMSSchema:
         referenced_spaces |= {s.space for s in self.spaces.values()}
         return referenced_spaces
 
+    def referenced_container(self) -> set[dm.ContainerId]:
+        referenced_containers = {
+            container for view in self.views.values() for container in view.referenced_containers()
+        }
+        referenced_containers |= set(self.containers.keys())
+        return referenced_containers
+
     def as_read_model(self) -> dm.DataModel[dm.View]:
         if self.data_model is None:
             raise ValueError("Data model is not defined")
@@ -828,263 +839,3 @@ class DMSSchema:
         }
         # If a container has a required property that is not used by the view, the view is not writable
         return not bool(required_properties - used_properties)
-
-
-@dataclass
-class PipelineSchema(DMSSchema):
-    transformations: TransformationWriteList = field(default_factory=lambda: TransformationWriteList([]))
-    databases: DatabaseWriteList = field(default_factory=lambda: DatabaseWriteList([]))
-    raw_tables: RawTableWriteList = field(default_factory=lambda: RawTableWriteList([]))
-
-    _FIELD_NAME_BY_RESOURCE_TYPE: ClassVar[dict[str, str]] = {
-        **DMSSchema._FIELD_NAME_BY_RESOURCE_TYPE,
-        "raw": "raw_tables",
-    }
-
-    def __post_init__(self):
-        existing_databases = {database.name for database in self.databases}
-        table_database = {table.database for table in self.raw_tables}
-        if missing := table_database - existing_databases:
-            self.databases.extend([DatabaseWrite(name=database) for database in missing])
-
-    @classmethod
-    def _read_directory(cls, directory: Path) -> tuple[dict[str, list[Any]], dict[str, list[Path]]]:
-        data, context = super()._read_directory(directory)
-        for yaml_file in directory.rglob("*.yaml"):
-            if yaml_file.parent.name in ("transformations", "raw"):
-                attr_name = cls._FIELD_NAME_BY_RESOURCE_TYPE.get(yaml_file.parent.name, yaml_file.parent.name)
-                data.setdefault(attr_name, [])
-                context.setdefault(attr_name, [])
-                try:
-                    loaded = yaml.safe_load(yaml_file.read_text())
-                except Exception as e:
-                    warnings.warn(
-                        FileTypeUnexpectedWarning(yaml_file, frozenset([".yaml", ".yml"]), str(e)), stacklevel=2
-                    )
-                    continue
-                if isinstance(loaded, list):
-                    data[attr_name].extend(loaded)
-                    context[attr_name].extend([yaml_file] * len(loaded))
-                else:
-                    data[attr_name].append(loaded)
-                    context[attr_name].append(yaml_file)
-        return data, context
-
-    def to_directory(
-        self,
-        directory: str | Path,
-        exclude: set[str] | None = None,
-        new_line: str | None = "\n",
-        encoding: str | None = "utf-8",
-    ) -> None:
-        super().to_directory(directory, exclude)
-        exclude_set = exclude or set()
-        path_dir = Path(directory)
-        if "transformations" not in exclude_set and self.transformations:
-            transformation_dir = path_dir / "transformations"
-            transformation_dir.mkdir(exist_ok=True, parents=True)
-            for transformation in self.transformations:
-                (transformation_dir / f"{transformation.external_id}.yaml").write_text(
-                    transformation.dump_yaml(), newline=new_line, encoding=encoding
-                )
-        if "raw" not in exclude_set and self.raw_tables:
-            # The RAW Databases are not written to file. This is because cognite-toolkit expects the RAW databases
-            # to be in the same file as the RAW tables.
-            raw_dir = path_dir / "raw"
-            raw_dir.mkdir(exist_ok=True, parents=True)
-            for raw_table in self.raw_tables:
-                (raw_dir / f"{raw_table.name}.yaml").write_text(
-                    raw_table.dump_yaml(), newline=new_line, encoding=encoding
-                )
-
-    def to_zip(self, zip_file: str | Path, exclude: set[str] | None = None) -> None:
-        super().to_zip(zip_file, exclude)
-        exclude_set = exclude or set()
-        with zipfile.ZipFile(zip_file, "a") as zip_ref:
-            if "transformations" not in exclude_set:
-                for transformation in self.transformations:
-                    zip_ref.writestr(f"transformations/{transformation.external_id}.yaml", transformation.dump_yaml())
-            if "raw" not in exclude_set:
-                # The RAW Databases are not written to file. This is because cognite-toolkit expects the RAW databases
-                # to be in the same file as the RAW tables.
-                for raw_table in self.raw_tables:
-                    zip_ref.writestr(f"raw/{raw_table.name}.yaml", raw_table.dump_yaml())
-
-    @classmethod
-    def _read_zip(cls, zip_file: Path) -> tuple[dict[str, list[Any]], dict[str, list[Path]]]:
-        data, context = super()._read_zip(zip_file)
-        with zipfile.ZipFile(zip_file, "r") as zip_ref:
-            for file_info in zip_ref.infolist():
-                if file_info.filename.endswith(".yaml"):
-                    if "/" not in file_info.filename:
-                        continue
-                    filepath = Path(file_info.filename)
-                    if (parent := filepath.parent.name) in ("transformations", "raw"):
-                        attr_name = cls._FIELD_NAME_BY_RESOURCE_TYPE.get(parent, parent)
-                        data.setdefault(attr_name, [])
-                        context.setdefault(attr_name, [])
-                        try:
-                            loaded = yaml.safe_load(zip_ref.read(file_info).decode())
-                        except Exception as e:
-                            warnings.warn(
-                                FileTypeUnexpectedWarning(filepath, frozenset([".yaml", ".yml"]), str(e)), stacklevel=2
-                            )
-                            continue
-                        if isinstance(loaded, list):
-                            data[attr_name].extend(loaded)
-                            context[attr_name].extend([filepath] * len(loaded))
-                        else:
-                            data[attr_name].append(loaded)
-                            context[attr_name].append(filepath)
-        return data, context
-
-    @classmethod
-    def from_dms(cls, schema: DMSSchema, instance_space: str | None = None) -> "PipelineSchema":
-        if not schema.data_model:
-            raise ValueError("PipelineSchema must contain at least one data model")
-        first_data_model = schema.data_model
-        # The database name is limited to 32 characters
-        database_name = first_data_model.external_id[:32]
-        instance_space = instance_space or first_data_model.space
-        database = DatabaseWrite(name=database_name)
-        parent_views = {parent for view in schema.views.values() for parent in view.implements or []}
-        container_by_id = schema.containers.copy()
-
-        transformations = TransformationWriteList([])
-        raw_tables = RawTableWriteList([])
-        for view in schema.views.values():
-            if view.as_id() in parent_views:
-                # Skipping parents as they do not have their own data
-                continue
-            mapped_properties = {
-                prop_name: prop
-                for prop_name, prop in (view.properties or {}).items()
-                if isinstance(prop, dm.MappedPropertyApply)
-            }
-            if mapped_properties:
-                view_table = RawTableWrite(name=f"{view.external_id}Properties", database=database_name)
-                raw_tables.append(view_table)
-                transformation = cls._create_property_transformation(
-                    mapped_properties, view, view_table, container_by_id, instance_space
-                )
-                transformations.append(transformation)
-            connection_properties = {
-                prop_name: prop
-                for prop_name, prop in (view.properties or {}).items()
-                if isinstance(prop, dm.EdgeConnectionApply)
-            }
-            for prop_name, connection_property in connection_properties.items():
-                view_table = RawTableWrite(name=f"{view.external_id}.{prop_name}Edge", database=database_name)
-                raw_tables.append(view_table)
-                transformation = cls._create_edge_transformation(connection_property, view, view_table, instance_space)
-                transformations.append(transformation)
-
-        return cls(
-            spaces=schema.spaces,
-            data_model=schema.data_model,
-            views=schema.views,
-            containers=schema.containers,
-            transformations=transformations,
-            databases=DatabaseWriteList([database]),
-            raw_tables=raw_tables,
-        )
-
-    @classmethod
-    def _create_property_transformation(
-        cls,
-        properties: dict[str, dm.MappedPropertyApply],
-        view: ViewApply,
-        table: RawTableWrite,
-        container_by_id: dict[dm.ContainerId, dm.ContainerApply],
-        instance_space: str,
-    ) -> TransformationWrite:
-        mapping_mode = {
-            "version": 1,
-            "sourceType": "raw",
-            # 'mappings' is set here and overwritten further down to ensure the correct order
-            "mappings": [],
-            "sourceLevel1": table.database,
-            "sourceLevel2": table.name,
-        }
-        mappings = [
-            {"from": "externalId", "to": "externalId", "asType": "STRING"},
-        ]
-        select_rows = ["cast(`externalId` as STRING) as externalId"]
-        for prop_name, prop in properties.items():
-            container = container_by_id.get(prop.container)
-            if container is not None:
-                dms_type = container.properties[prop.container_property_identifier].type._type
-                if dms_type in _DATA_TYPE_BY_DMS_TYPE:
-                    sql_type = _DATA_TYPE_BY_DMS_TYPE[dms_type].sql
-                else:
-                    warnings.warn(
-                        f"Unknown DMS type '{dms_type}' for property '{prop_name}'", RuntimeWarning, stacklevel=2
-                    )
-                    sql_type = "STRING"
-            else:
-                sql_type = "STRING"
-            select_rows.append(f"cast(`{prop_name}` as {sql_type}) as {prop_name}")
-            mappings.append({"from": prop_name, "to": prop_name, "asType": sql_type})
-        mapping_mode["mappings"] = mappings
-        select = ",\n  ".join(select_rows)
-
-        return TransformationWrite(
-            external_id=f"{table.name}Transformation",
-            name=f"{table.name}Transformation",
-            ignore_null_fields=True,
-            destination=Nodes(
-                view=ViewInfo(view.space, view.external_id, view.version),
-                instance_space=instance_space,
-            ),
-            conflict_mode="upsert",
-            query=f"""/* MAPPING_MODE_ENABLED: true */
-/* {json.dumps(mapping_mode)} */
-select
-  {select}
-from
-  `{table.database}`.`{table.name}`;
-""",
-        )
-
-    @classmethod
-    def _create_edge_transformation(
-        cls, property_: dm.EdgeConnectionApply, view: ViewApply, table: RawTableWrite, instance_space: str
-    ) -> TransformationWrite:
-        start, end = view.external_id, property_.source.external_id
-        if property_.direction == "inwards":
-            start, end = end, start
-        mapping_mode = {
-            "version": 1,
-            "sourceType": "raw",
-            "mappings": [
-                {"from": "externalId", "to": "externalId", "asType": "STRING"},
-                {"from": start, "to": "startNode", "asType": "STRUCT<`space`:STRING, `externalId`:STRING>"},
-                {"from": end, "to": "endNode", "asType": "STRUCT<`space`:STRING, `externalId`:STRING>"},
-            ],
-            "sourceLevel1": table.database,
-            "sourceLevel2": table.name,
-        }
-        select_rows = [
-            "cast(`externalId` as STRING) as externalId",
-            f"node_reference('{instance_space}', `{start}`) as startNode",
-            f"node_reference('{instance_space}', `{end}`) as endNode",
-        ]
-        select = ",\n  ".join(select_rows)
-
-        return TransformationWrite(
-            external_id=f"{table.name}Transformation",
-            name=f"{table.name}Transformation",
-            ignore_null_fields=True,
-            destination=Edges(
-                instance_space=instance_space,
-                edge_type=EdgeType(space=property_.type.space, external_id=property_.type.external_id),
-            ),
-            conflict_mode="upsert",
-            query=f"""/* MAPPING_MODE_ENABLED: true */
-/* {json.dumps(mapping_mode)} */
-select
-  {select}
-from
-  `{table.database}`.`{table.name}`;
-""",
-        )

cognite/neat/_graph/_shared.py

@@ -4,31 +4,30 @@ MIMETypes: TypeAlias = Literal[
     "application/rdf+xml", "text/turtle", "application/n-triple", "application/n-quads", "application/trig"
 ]
 
+RDFTypes: TypeAlias = Literal["xml", "rdf", "owl", "n3", "ttl", "turtle", "nt", "nq", "nquads", "trig"]
 
-def rdflib_to_mime_types(rdflib_format: str) -> str | None:
+
+def rdflib_to_oxi_type(rdflib_format: str) -> str | None:
     """Convert an RDFlib format to a MIME type.
 
     Args:
         rdflib_format (str): The RDFlib format.
 
     Returns:
-        MIMETypes: The MIME type.
-
-    !!! note
-        This will be replaced once new version of oxrdflib is released.
+        Oxi format used to trigger correct plugging in rdflib
 
     """
 
     mapping = {
-        "xml": "application/rdf+xml",
-        "rdf": "application/rdf+xml",
-        "owl": "application/rdf+xml",
-        "n3": "application/n-triple",
-        "ttl": "text/turtle",
-        "turtle": "text/turtle",
-        "nt": "application/n-triple",
-        "nq": "application/n-quads",
-        "nquads": "application/n-quads",
-        "trig": "application/trig",
+        "xml": "ox-xml",
+        "rdf": "ox-xml",
+        "owl": "ox-xml",
+        "n3": "ox-n3",
+        "ttl": "ox-ttl",
+        "turtle": "ox-turtle",
+        "nt": "ox-nt",
+        "nq": "ox-nq",
+        "nquads": "ox-nquads",
+        "trig": "ox-trig",
     }
     return mapping.get(rdflib_format, None)