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

cognite/neat/_session/_prepare.py

@@ -1,26 +1,44 @@
+import copy
 from collections.abc import Collection
 from datetime import datetime, timezone
-from typing import Literal
+from typing import Literal, cast
 
 from cognite.client.data_classes.data_modeling import DataModelIdentifier
 from rdflib import URIRef
 
+from cognite.neat._client import NeatClient
+from cognite.neat._constants import DEFAULT_NAMESPACE
+from cognite.neat._graph.transformers import RelationshipToSchemaTransformer
 from cognite.neat._graph.transformers._rdfpath import MakeConnectionOnExactMatch
-from cognite.neat._rules._shared import ReadRules
+from cognite.neat._rules._shared import InputRules, ReadRules
+from cognite.neat._rules.importers import DMSImporter
+from cognite.neat._rules.models import DMSRules
 from cognite.neat._rules.models.information._rules_input import InformationInputRules
-from cognite.neat._rules.transformers import ReduceCogniteModel, ToCompliantEntities, ToExtension
+from cognite.neat._rules.transformers import (
+    PrefixEntities,
+    ReduceCogniteModel,
+    ToCompliantEntities,
+    ToExtension,
+    VerifyDMSRules,
+)
+from cognite.neat._store._provenance import Agent as ProvenanceAgent
 from cognite.neat._store._provenance import Change
 
 from ._state import SessionState
 from .exceptions import NeatSessionError, session_class_wrapper
 
+try:
+    from rich import print
+except ImportError:
+    ...
+
 
 @session_class_wrapper
 class PrepareAPI:
-    def __init__(self, state: SessionState, verbose: bool) -> None:
+    def __init__(self, client: NeatClient | None, state: SessionState, verbose: bool) -> None:
         self._state = state
         self._verbose = verbose
-        self.data_model = DataModelPrepareAPI(state, verbose)
+        self.data_model = DataModelPrepareAPI(client, state, verbose)
         self.instances = InstancePrepareAPI(state, verbose)
 
 
@@ -94,34 +112,77 @@ class InstancePrepareAPI:
             raise NeatSessionError(f"Property {property_} is not defined for type {type_}. Cannot make connection")
         return type_uri[0], property_uri[0]
 
+    def relationships_as_connections(self, limit: int = 1) -> None:
+        """This assumes that you have read a classic CDF knowledge graph including relationships.
+
+        This transformer analyzes the relationships in the graph and modifies them to be part of the schema
+        for Assets, Events, Files, Sequences, and TimeSeries. Relationships without any properties
+        are replaced by a simple relationship between the source and target nodes. Relationships with
+        properties are replaced by a schema that contains the properties as attributes.
+
+        Args:
+            limit: The minimum number of relationships that need to be present for it
+                to be converted into a schema. Default is 1.
+
+        """
+        transformer = RelationshipToSchemaTransformer(limit=limit)
+        self._state.instances.store.transform(transformer)
+
 
 @session_class_wrapper
 class DataModelPrepareAPI:
-    def __init__(self, state: SessionState, verbose: bool) -> None:
+    def __init__(self, client: NeatClient | None, state: SessionState, verbose: bool) -> None:
+        self._client = client
         self._state = state
         self._verbose = verbose
 
     def cdf_compliant_external_ids(self) -> None:
         """Convert data model component external ids to CDF compliant entities."""
-        if input := self._state.data_model.last_info_unverified_rule:
-            source_id, rules = input
+        source_id, rules = self._state.data_model.last_info_unverified_rule
+
+        start = datetime.now(timezone.utc)
+        transformer = ToCompliantEntities()
+        output: ReadRules[InformationInputRules] = transformer.transform(rules)
+        end = datetime.now(timezone.utc)
+
+        change = Change.from_rules_activity(
+            output,
+            transformer.agent,
+            start,
+            end,
+            "Converted external ids to CDF compliant entities",
+            self._state.data_model.provenance.source_entity(source_id)
+            or self._state.data_model.provenance.target_entity(source_id),
+        )
 
-            start = datetime.now(timezone.utc)
-            transformer = ToCompliantEntities()
-            output: ReadRules[InformationInputRules] = transformer.transform(rules)
-            end = datetime.now(timezone.utc)
+        self._state.data_model.write(output, change)
 
-            change = Change.from_rules_activity(
-                output,
-                transformer.agent,
-                start,
-                end,
-                "Converted external ids to CDF compliant entities",
-                self._state.data_model.provenance.source_entity(source_id)
-                or self._state.data_model.provenance.target_entity(source_id),
-            )
+    def prefix(self, prefix: str) -> None:
+        """Prefix all views in the data model with the given prefix.
 
-            self._state.data_model.write(output, change)
+        Args:
+            prefix: The prefix to add to the views in the data model.
+
+        """
+        source_id, rules = self._state.data_model.last_unverified_rule
+
+        start = datetime.now(timezone.utc)
+        transformer = PrefixEntities(prefix)
+        new_rules = cast(InputRules, copy.deepcopy(rules.get_rules()))
+        output = transformer.transform(new_rules)
+        end = datetime.now(timezone.utc)
+
+        change = Change.from_rules_activity(
+            output,
+            transformer.agent,
+            start,
+            end,
+            "Added prefix to the data model views",
+            self._state.data_model.provenance.source_entity(source_id)
+            or self._state.data_model.provenance.target_entity(source_id),
+        )
+
+        self._state.data_model.write(output, change)
 
     def to_enterprise(
         self,
@@ -185,7 +246,7 @@ class DataModelPrepareAPI:
         data_model_id: DataModelIdentifier,
         org_name: str = "My",
         mode: Literal["read", "write"] = "read",
-        dummy_property: str = "dummy",
+        dummy_property: str = "GUID",
     ) -> None:
         """Uses the current data model as a basis to create solution data model
 
@@ -235,6 +296,81 @@ class DataModelPrepareAPI:
 
             self._state.data_model.write(output.rules, change)
 
+    def to_data_product(
+        self,
+        data_model_id: DataModelIdentifier,
+        org_name: str = "",
+        include: Literal["same-space", "all"] = "same-space",
+    ) -> None:
+        """Uses the current data model as a basis to create data product data model.
+
+        A data product model is a data model that ONLY maps to containers and do not use implements. This is
+        typically used for defining the data in a data product.
+
+        Args:
+            data_model_id: The data product data model id that is being created.
+            org_name: Organization name to use for the views in the new data model.
+            include: The views to include in the data product data model. Can be either "same-space" or "all".
+                If you set same-space, only the views in the same space as the data model will be included.
+        """
+        source_id, rules = self._state.data_model.last_verified_dms_rules
+
+        dms_ref: DMSRules | None = None
+        view_ids, container_ids = rules.imported_views_and_containers_ids(include_model_views_with_no_properties=True)
+        if view_ids or container_ids:
+            if self._client is None:
+                raise NeatSessionError(
+                    "No client provided. You are referencing unknown views and containers in your data model, "
+                    "NEAT needs a client to lookup the definitions. "
+                    "Please set the client in the session, NeatSession(client=client)."
+                )
+            schema = self._client.schema.retrieve(list(view_ids), list(container_ids))
+
+            importer = DMSImporter(schema)
+            reference_rules = importer.to_rules().rules
+            if reference_rules is not None:
+                imported = VerifyDMSRules("continue").transform(reference_rules)
+                if dms_ref := imported.rules:
+                    rules = rules.model_copy(deep=True)
+                    if rules.containers is None:
+                        rules.containers = dms_ref.containers
+                    else:
+                        existing_containers = {c.container for c in rules.containers}
+                        rules.containers.extend(
+                            [c for c in dms_ref.containers or [] if c.container not in existing_containers]
+                        )
+                    existing_views = {v.view for v in rules.views}
+                    rules.views.extend([v for v in dms_ref.views if v.view not in existing_views])
+                    existing_properties = {(p.view, p.view_property) for p in rules.properties}
+                    rules.properties.extend(
+                        [p for p in dms_ref.properties if (p.view, p.view_property) not in existing_properties]
+                    )
+
+        start = datetime.now(timezone.utc)
+        transformer = ToExtension(
+            new_model_id=data_model_id,
+            org_name=org_name,
+            type_="data_product",
+            include=include,
+        )
+        output = transformer.transform(rules)
+        end = datetime.now(timezone.utc)
+
+        change = Change.from_rules_activity(
+            output,
+            transformer.agent,
+            start,
+            end,
+            (
+                f"Prepared data model {data_model_id} to be data product model "
+                f"on top of {rules.metadata.as_data_model_id()}"
+            ),
+            self._state.data_model.provenance.source_entity(source_id)
+            or self._state.data_model.provenance.target_entity(source_id),
+        )
+
+        self._state.data_model.write(output.rules, change)
+
     def reduce(self, drop: Collection[Literal["3D", "Annotation", "BaseViews"] | str]) -> None:
         """This is a special method that allow you to drop parts of the data model.
         This only applies to Cognite Data Models.
@@ -267,3 +403,62 @@ class DataModelPrepareAPI:
             )
 
             self._state.data_model.write(output.rules, change)
+
+    def include_referenced(self) -> None:
+        """Include referenced views and containers in the data model."""
+        start = datetime.now(timezone.utc)
+
+        source_id, rules = self._state.data_model.last_verified_dms_rules
+        view_ids, container_ids = rules.imported_views_and_containers_ids(include_model_views_with_no_properties=True)
+        if not (view_ids or container_ids):
+            print(
+                f"Data model {rules.metadata.as_data_model_id()} does not have any referenced views or containers."
+                f"that is not already included in the data model."
+            )
+            return
+        if self._client is None:
+            raise NeatSessionError(
+                "No client provided. You are referencing unknown views and containers in your data model, "
+                "NEAT needs a client to lookup the definitions. "
+                "Please set the client in the session, NeatSession(client=client)."
+            )
+        schema = self._client.schema.retrieve(list(view_ids), list(container_ids))
+        copy_ = rules.model_copy(deep=True)
+        copy_.metadata.version = f"{rules.metadata.version}_completed"
+        importer = DMSImporter(schema)
+        imported = importer.to_rules()
+        if imported.rules is None:
+            self._state.data_model.issue_lists.append(imported.issues)
+            raise NeatSessionError(
+                "Could not import the referenced views and containers. "
+                "See `neat.inspect.issues()` for more information."
+            )
+        verified = VerifyDMSRules("continue", post_validate=False).transform(imported.rules)
+        if verified.rules is None:
+            self._state.data_model.issue_lists.append(verified.issues)
+            raise NeatSessionError(
+                "Could not verify the referenced views and containers. "
+                "See `neat.inspect.issues()` for more information."
+            )
+        if copy_.containers is None:
+            copy_.containers = verified.rules.containers
+        else:
+            existing_containers = {c.container for c in copy_.containers}
+            copy_.containers.extend(
+                [c for c in verified.rules.containers or [] if c.container not in existing_containers]
+            )
+        existing_views = {v.view for v in copy_.views}
+        copy_.views.extend([v for v in verified.rules.views if v.view not in existing_views])
+        end = datetime.now(timezone.utc)
+
+        change = Change.from_rules_activity(
+            copy_,
+            ProvenanceAgent(id_=DEFAULT_NAMESPACE["agent/"]),
+            start,
+            end,
+            (f"Included referenced views and containers in the data model {rules.metadata.as_data_model_id()}"),
+            self._state.data_model.provenance.source_entity(source_id)
+            or self._state.data_model.provenance.target_entity(source_id),
+        )
+
+        self._state.data_model.write(copy_, change)

cognite/neat/_session/_read.py

@@ -3,9 +3,9 @@ from datetime import datetime, timezone
 from pathlib import Path
 from typing import Any, Literal
 
-from cognite.client import CogniteClient
 from cognite.client.data_classes.data_modeling import DataModelId, DataModelIdentifier
 
+from cognite.neat._client import NeatClient
 from cognite.neat._constants import COGNITE_SPACES
 from cognite.neat._graph import examples as instances_examples
 from cognite.neat._graph import extractors
@@ -27,7 +27,7 @@ from .exceptions import NeatSessionError, session_class_wrapper
 
 @session_class_wrapper
 class ReadAPI:
-    def __init__(self, state: SessionState, client: CogniteClient | None, verbose: bool) -> None:
+    def __init__(self, state: SessionState, client: NeatClient | None, verbose: bool) -> None:
         self._state = state
         self._verbose = verbose
         self.cdf = CDFReadAPI(state, client, verbose)
@@ -39,7 +39,7 @@ class ReadAPI:
 
 @session_class_wrapper
 class BaseReadAPI:
-    def __init__(self, state: SessionState, client: CogniteClient | None, verbose: bool) -> None:
+    def __init__(self, state: SessionState, client: NeatClient | None, verbose: bool) -> None:
         self._state = state
         self._verbose = verbose
         self._client = client
@@ -67,12 +67,12 @@ class BaseReadAPI:
 
 @session_class_wrapper
 class CDFReadAPI(BaseReadAPI):
-    def __init__(self, state: SessionState, client: CogniteClient | None, verbose: bool) -> None:
+    def __init__(self, state: SessionState, client: NeatClient | None, verbose: bool) -> None:
         super().__init__(state, client, verbose)
         self.classic = CDFClassicAPI(state, client, verbose)
 
     @property
-    def _get_client(self) -> CogniteClient:
+    def _get_client(self) -> NeatClient:
         if self._client is None:
             raise NeatValueError("No client provided. Please provide a client to read a data model.")
         return self._client
@@ -113,16 +113,53 @@ class CDFReadAPI(BaseReadAPI):
 @session_class_wrapper
 class CDFClassicAPI(BaseReadAPI):
     @property
-    def _get_client(self) -> CogniteClient:
+    def _get_client(self) -> NeatClient:
         if self._client is None:
             raise ValueError("No client provided. Please provide a client to read a data model.")
         return self._client
 
-    def assets(self, root_asset_external_id: str) -> None:
-        extractor = extractors.AssetsExtractor.from_hierarchy(self._get_client, root_asset_external_id)
+    def graph(self, root_asset_external_id: str) -> None:
+        """Reads the classic knowledge graph from CDF.
+
+        The Classic Graph consists of the following core resource type.
+
+        Classic Node CDF Resources:
+         - Assets
+         - TimeSeries
+         - Sequences
+         - Events
+         - Files
+
+        All the classic node CDF resources can have one or more connections to one or more assets. This
+        will match a direct relationship in the data modeling of CDF.
+
+        In addition, you have relationships between the classic node CDF resources. This matches an edge
+        in the data modeling of CDF.
+
+        Finally, you have labels and data sets that to organize the graph. In which data sets have a similar,
+        but different, role as a space in data modeling. While labels can be compared to node types in data modeling,
+        used to quickly filter and find nodes/edges.
+
+        This extractor will extract the classic CDF graph into Neat starting from either a data set or a root asset.
+
+        It works as follows:
+
+        1. Extract all core nodes (assets, time series, sequences, events, files) filtered by the given data set or
+           root asset.
+        2. Extract all relationships starting from any of the extracted core nodes.
+        3. Extract all core nodes that are targets of the relationships that are not already extracted.
+        4. Extract all labels that are connected to the extracted core nodes/relationships.
+        5. Extract all data sets that are connected to the extracted core nodes/relationships.
+
+        Args:
+            root_asset_external_id: The external id of the root asset
+
+        """
+        extractor = extractors.ClassicGraphExtractor(self._get_client, root_asset_external_id=root_asset_external_id)
+
         self._state.instances.store.write(extractor)
         if self._verbose:
-            print(f"Asset hierarchy {root_asset_external_id} read successfully")
+            print(f"Classic Graph {root_asset_external_id} read successfully")
 
 
 @session_class_wrapper
@@ -145,7 +182,7 @@ class ExcelReadAPI(BaseReadAPI):
                 description=f"Excel file {reader!s} read as unverified data model",
             )
             self._store_rules(input_rules, change)
-
+        self._state.data_model.issue_lists.append(input_rules.issues)
         return input_rules.issues
 
 
@@ -176,7 +213,7 @@ class YamlReadAPI(BaseReadAPI):
                         "NEAT needs a client to lookup the container definitions. "
                         "Please set the client in the session, NeatSession(client=client)."
                     )
-                system_containers = self._client.data_modeling.containers.retrieve(system_container_ids)
+                system_containers = self._client.loaders.containers.retrieve(system_container_ids)
                 dms_importer.update_referenced_containers(system_containers)
 
             importer = dms_importer
@@ -222,7 +259,7 @@ class CSVReadAPI(BaseReadAPI):
 
 @session_class_wrapper
 class RDFReadAPI(BaseReadAPI):
-    def __init__(self, state: SessionState, client: CogniteClient | None, verbose: bool) -> None:
+    def __init__(self, state: SessionState, client: NeatClient | None, verbose: bool) -> None:
         super().__init__(state, client, verbose)
         self.examples = RDFExamples(state)
 

cognite/neat/_session/_to.py

@@ -1,9 +1,9 @@
 from pathlib import Path
 from typing import Any, Literal, overload
 
-from cognite.client import CogniteClient
 from cognite.client.data_classes.data_modeling import SpaceApply
 
+from cognite.neat._client import NeatClient
 from cognite.neat._graph import loaders
 from cognite.neat._issues import IssueList, catch_warnings
 from cognite.neat._rules import exporters
@@ -17,7 +17,7 @@ from .exceptions import NeatSessionError, session_class_wrapper
 
 @session_class_wrapper
 class ToAPI:
-    def __init__(self, state: SessionState, client: CogniteClient | None, verbose: bool) -> None:
+    def __init__(self, state: SessionState, client: NeatClient | None, verbose: bool) -> None:
         self._state = state
         self._verbose = verbose
         self.cdf = CDFToAPI(state, client, verbose)
@@ -75,7 +75,7 @@ class ToAPI:
 
 @session_class_wrapper
 class CDFToAPI:
-    def __init__(self, state: SessionState, client: CogniteClient | None, verbose: bool) -> None:
+    def __init__(self, state: SessionState, client: NeatClient | None, verbose: bool) -> None:
         self._client = client
         self._state = state
         self._verbose = verbose

cognite/neat/_store/_base.py

@@ -7,10 +7,11 @@ from typing import cast
 
 import pandas as pd
 from pandas import Index
-from rdflib import Graph, Namespace, URIRef
+from rdflib import Dataset, Namespace, URIRef
 from rdflib.plugins.stores.sparqlstore import SPARQLUpdateStore
 
 from cognite.neat._constants import DEFAULT_NAMESPACE
+from cognite.neat._graph._shared import rdflib_to_oxi_type
 from cognite.neat._graph.extractors import RdfFileExtractor, TripleExtractors
 from cognite.neat._graph.queries import Queries
 from cognite.neat._graph.transformers import Transformers
@@ -42,7 +43,7 @@ class NeatGraphStore:
 
     def __init__(
         self,
-        graph: Graph,
+        graph: Dataset,
         rules: InformationRules | None = None,
     ):
         self.rules: InformationRules | None = None
@@ -109,7 +110,7 @@ class NeatGraphStore:
 
     @classmethod
     def from_memory_store(cls, rules: InformationRules | None = None) -> "Self":
-        return cls(Graph(identifier=DEFAULT_NAMESPACE), rules)
+        return cls(Dataset(), rules)
 
     @classmethod
     def from_sparql_store(
@@ -127,7 +128,7 @@ class NeatGraphStore:
             postAsEncoded=False,
             autocommit=False,
         )
-        graph = Graph(store=store, identifier=DEFAULT_NAMESPACE)
+        graph = Dataset(store=store)
         return cls(graph, rules)
 
     @classmethod
@@ -150,9 +151,8 @@ class NeatGraphStore:
         else:
             raise Exception("Error initializing Oxigraph store")
 
-        graph = Graph(
+        graph = Dataset(
             store=oxrdflib.OxigraphStore(store=oxi_store),
-            identifier=DEFAULT_NAMESPACE,
         )
 
         return cls(graph, rules)
@@ -162,7 +162,7 @@ class NeatGraphStore:
         success = True
 
         if isinstance(extractor, RdfFileExtractor) and not extractor.issue_list.has_errors:
-            self._parse_file(extractor.filepath, cast(str, extractor.mime_type), extractor.base_uri)
+            self._parse_file(extractor.filepath, cast(str, extractor.format), extractor.base_uri)
         elif isinstance(extractor, RdfFileExtractor):
             success = False
             issue_text = "\n".join([issue.as_message() for issue in extractor.issue_list])
@@ -244,33 +244,36 @@ class NeatGraphStore:
     def _parse_file(
         self,
         filepath: Path,
-        mime_type: str = "application/rdf+xml",
+        format: str = "turtle",
         base_uri: URIRef | None = None,
     ) -> None:
         """Imports graph data from file.
 
         Args:
             filepath : File path to file containing graph data, by default None
-            mime_type : MIME type of graph data, by default "application/rdf+xml"
-            base_uri : Add base IRI to graph, by default True
+            format : rdflib format file containing RDF graph, by default "turtle"
+            base_uri : base URI to add to graph in case of relative URIs, by default None
+
+        !!! note "Oxigraph store"
+            By default we are using non-transactional mode for parsing RDF files.
+            This gives us a significant performance boost when importing large RDF files.
+            Underhood of rdflib we are triggering oxrdflib plugin which in respect
+            calls `bulk_load` method from oxigraph store. See more at:
+            https://pyoxigraph.readthedocs.io/en/stable/store.html#pyoxigraph.Store.bulk_load
         """
 
         # Oxigraph store, do not want to type hint this as it is an optional dependency
         if type(self.graph.store).__name__ == "OxigraphStore":
-
-            def parse_to_oxi_store():
-                local_import("pyoxigraph", "oxi")
-                import pyoxigraph
-
-                cast(pyoxigraph.Store, self.graph.store._store).bulk_load(
-                    str(filepath),
-                    mime_type,
-                    base_iri=base_uri,
-                    to_graph=pyoxigraph.NamedNode(self.graph.identifier),
-                )
-                cast(pyoxigraph.Store, self.graph.store._store).optimize()
-
-            parse_to_oxi_store()
+            local_import("pyoxigraph", "oxi")
+
+            # this is necessary to trigger rdflib oxigraph plugin
+            self.graph.parse(
+                filepath,
+                format=rdflib_to_oxi_type(format),
+                transactional=False,
+                publicID=base_uri,
+            )
+            self.graph.store._store.optimize()  # type: ignore[attr-defined]
 
         # All other stores
         else:

cognite/neat/_utils/rdf_.py

@@ -173,7 +173,7 @@ def get_inheritance_path(child: Any, child_parent: dict[Any, list[Any]]) -> list
     return path
 
 
-def add_triples_in_batch(graph: Graph, triples: Iterable[Triple], batch_size: int = 10_000):
+def add_triples_in_batch(graph: Graph, triples: Iterable[Triple], batch_size: int = 10_000) -> None:
     """Adds triples to the graph store in batches.
 
     Args:
@@ -204,3 +204,30 @@ def add_triples_in_batch(graph: Graph, triples: Iterable[Triple], batch_size: in
         check_commit()
 
     check_commit(force_commit=True)
+
+
+def remove_instance_ids_in_batch(graph: Graph, instance_ids: Iterable[URIRef], batch_size: int = 1_000) -> None:
+    """Removes all triples related to the given instances in the graph store in batches.
+
+    Args:
+        graph: The graph store to remove triples from
+        instance_ids:  list of instances to remove triples from
+        batch_size:  Batch size of triples per commit, by default 10_000
+
+    """
+    batch_count = 0
+
+    def check_commit(force_commit: bool = False):
+        """Commit nodes to the graph if batch counter is reached or if force_commit is True"""
+        nonlocal batch_count
+        batch_count += 1
+        if force_commit or batch_count >= batch_size:
+            graph.commit()
+            batch_count = 0
+            return
+
+    for instance_id in instance_ids:
+        graph.remove((instance_id, None, None))
+        check_commit()
+
+    check_commit(force_commit=True)

cognite/neat/_version.py

@@ -1,2 +1,2 @@
-__version__ = "0.98.0"
+__version__ = "0.99.0"
 __engine__ = "^1.0.3"

cognite/neat/_workflows/steps/lib/current/rules_exporter.py

@@ -2,6 +2,7 @@ import time
 from pathlib import Path
 from typing import ClassVar, Literal, cast
 
+from cognite.neat._client import NeatClient
 from cognite.neat._issues.errors import WorkflowStepNotInitializedError
 from cognite.neat._rules import exporters
 from cognite.neat._rules._shared import DMSRules, InformationRules, VerifiedRules
@@ -100,7 +101,7 @@ class DeleteDataModelFromCDF(Step):
 
         report_lines = ["# Data Model Deletion from CDF\n\n"]
         errors = []
-        for result in dms_exporter.delete_from_cdf(rules=dms_rules, client=cdf_client, dry_run=dry_run):
+        for result in dms_exporter.delete_from_cdf(rules=dms_rules, client=NeatClient(cdf_client), dry_run=dry_run):
             report_lines.append(str(result))
             errors.extend(result.error_messages)
 
@@ -220,7 +221,9 @@ class RulesToDMS(Step):
 
         report_lines = ["# DMS Schema Export to CDF\n\n"]
         errors = []
-        for result in dms_exporter.export_to_cdf_iterable(rules=dms_rules, client=cdf_client, dry_run=dry_run):
+        for result in dms_exporter.export_to_cdf_iterable(
+            rules=dms_rules, client=NeatClient(cdf_client), dry_run=dry_run
+        ):
             report_lines.append(str(result))
             errors.extend(result.error_messages)
 
@@ -584,7 +587,9 @@ class RulesToCDFTransformations(Step):
 
         report_lines = ["# DMS Schema Export to CDF\n\n"]
         errors = []
-        for result in dms_exporter.export_to_cdf_iterable(rules=dms_rules, client=cdf_client, dry_run=dry_run):
+        for result in dms_exporter.export_to_cdf_iterable(
+            rules=dms_rules, client=NeatClient(cdf_client), dry_run=dry_run
+        ):
             report_lines.append(str(result))
             errors.extend(result.error_messages)
 

cognite/neat/_workflows/steps/lib/current/rules_importer.py

@@ -5,6 +5,7 @@ from typing import ClassVar
 from cognite.client import CogniteClient
 from cognite.client.data_classes.data_modeling import DataModelId
 
+from cognite.neat._client import NeatClient
 from cognite.neat._issues.errors import WorkflowStepNotInitializedError
 from cognite.neat._issues.formatters import FORMATTER_BY_NAME
 from cognite.neat._rules import importers
@@ -299,7 +300,9 @@ class DMSToRules(Step):
                 return FlowMessage(error_text=error_text, step_execution_status=StepExecutionStatus.ABORT_AND_FAIL)
             ref_model_id = ref_model.as_id()
 
-        dms_importer = importers.DMSImporter.from_data_model_id(cdf_client, datamodel_entity.as_id(), ref_model_id)
+        dms_importer = importers.DMSImporter.from_data_model_id(
+            NeatClient(cdf_client), datamodel_entity.as_id(), ref_model_id
+        )
 
         # if role is None, it will be inferred from the rules file
         role = self.configs.get("Role")