cognite-neat 0.75.8__py3-none-any.whl → 0.76.0__py3-none-any.whl

cognite/neat/config.py

@@ -1,39 +1,169 @@
+import json
+import logging
+import os
 import shutil
+import sys
 from pathlib import Path
+from typing import Literal, cast
+
+import yaml
+from pydantic import BaseModel, Field
+from yaml import safe_load
 
 from cognite.neat.constants import EXAMPLE_GRAPHS, EXAMPLE_RULES, EXAMPLE_WORKFLOWS
+from cognite.neat.utils.cdf import InteractiveCogniteClient, ServiceCogniteClient
+
+if sys.version_info >= (3, 11):
+    from enum import StrEnum
+    from typing import Self
+else:
+    from backports.strenum import StrEnum
+    from typing_extensions import Self
+
+LOG_FORMAT = "%(asctime)s.%(msecs)03d %(levelname)-8s %(message)s"
+LOG_DATE_FORMAT = "%Y-%m-%d %H:%M:%S"
+
+
+class RulesStoreType(StrEnum):
+    """Rules Store type"""
+
+    CDF = "cdf"
+    FILE = "file"
+    URL = "url"
+    GOOGLE_SHEET = "google_sheet"
+
+
+class WorkflowsStoreType(StrEnum):
+    """Workflows Store type"""
+
+    CDF = "cdf"
+    FILE = "file"
+    URL = "url"
+
+
+class Config(BaseModel):
+    workflows_store_type: WorkflowsStoreType = WorkflowsStoreType.FILE
+    data_store_path: Path = Field(default_factory=lambda: Path.cwd() / "data")
+
+    workflow_downloader_filter: list[str] | None = Field(
+        description="List of workflow names+tags to filter on when downloading workflows from CDF. "
+        "Example name:workflow_name=version,tag:tag_name",
+        default=None,
+    )
+
+    cdf_client: InteractiveCogniteClient | ServiceCogniteClient = ServiceCogniteClient()
+    cdf_default_dataset_id: int = 0
+    load_examples: bool = True
+
+    log_level: Literal["ERROR", "WARNING", "INFO", "DEBUG"] = "INFO"
+    log_format: str = LOG_FORMAT
+    download_workflows_from_cdf: bool = Field(
+        default=False, description="Downloads all workflows from CDF automatically and stores them locally"
+    )
+
+    stop_on_error: bool = False
 
+    @property
+    def _dir_suffix(self) -> str:
+        is_test_running = "pytest" in sys.modules
+        if is_test_running:
+            # Todo change the below to f"-{os.getpid()}" when all tests supports parallel execution.
+            return ""
+        return ""
 
-def copy_examples_to_directory(target_data_dir: Path):
+    @property
+    def rules_store_path(self) -> Path:
+        return self.data_store_path / f"rules{self._dir_suffix}"
+
+    @property
+    def workflows_store_path(self) -> Path:
+        return self.data_store_path / f"workflows{self._dir_suffix}"
+
+    @property
+    def source_graph_path(self) -> Path:
+        return self.data_store_path / f"source-graphs{self._dir_suffix}"
+
+    @property
+    def staging_path(self) -> Path:
+        return self.data_store_path / f"staging{self._dir_suffix}"
+
+    @classmethod
+    def from_yaml(cls, filepath: Path) -> Self:
+        return cls(**safe_load(filepath.read_text()))
+
+    def to_yaml(self, filepath: Path):
+        # Parse as json to avoid Path and Enum objects
+        dump = json.loads(self.json())
+
+        with filepath.open("w") as f:
+            yaml.safe_dump(dump, f)
+
+    @classmethod
+    def from_env(cls) -> Self:
+        missing = "Missing"
+        cdf_config = ServiceCogniteClient(
+            project=os.environ.get("NEAT_CDF_PROJECT", missing),
+            client_id=os.environ.get("NEAT_CDF_CLIENT_ID", missing),
+            client_secret=os.environ.get("NEAT_CDF_CLIENT_SECRET", missing),
+            base_url=os.environ.get("NEAT_CDF_BASE_URL", missing),
+            token_url=os.environ.get("NEAT_CDF_TOKEN_URL", missing),
+            scopes=[os.environ.get("NEAT_CDF_SCOPES", missing)],
+            timeout=int(os.environ.get("NEAT_CDF_CLIENT_TIMEOUT", "60")),
+            max_workers=int(os.environ.get("NEAT_CDF_CLIENT_MAX_WORKERS", "3")),
+        )
+
+        if workflow_downloader_filter_value := os.environ.get("NEAT_WORKFLOW_DOWNLOADER_FILTER", None):
+            workflow_downloader_filter = workflow_downloader_filter_value.split(",")
+        else:
+            workflow_downloader_filter = None
+
+        return cls(
+            cdf_client=cdf_config,
+            workflows_store_type=os.environ.get(  # type: ignore[arg-type]
+                "NEAT_WORKFLOWS_STORE_TYPE", WorkflowsStoreType.FILE
+            ),
+            data_store_path=Path(os.environ.get("NEAT_DATA_PATH", "/app/data")),
+            cdf_default_dataset_id=int(os.environ.get("NEAT_CDF_DEFAULT_DATASET_ID", 6476640149881990)),
+            log_level=cast(Literal["ERROR", "WARNING", "INFO", "DEBUG"], os.environ.get("NEAT_LOG_LEVEL", "INFO")),
+            workflow_downloader_filter=workflow_downloader_filter,
+            load_examples=bool(os.environ.get("NEAT_LOAD_EXAMPLES", True) in ["True", "true", "1"]),
+        )
+
+
+def copy_examples_to_directory(config: Config):
     """
     Copier over all the examples to the target_data_directory,
     without overwriting
 
     Args:
         target_data_dir : The target directory
+        suffix : The suffix to add to the directory names
 
     """
 
-    print(f"Copying examples into {target_data_dir}")
-    _copy_examples(EXAMPLE_RULES, target_data_dir / "rules")
-    _copy_examples(EXAMPLE_GRAPHS, target_data_dir / "source-graphs")
-    _copy_examples(EXAMPLE_WORKFLOWS, target_data_dir / "workflows")
-    (target_data_dir / "staging").mkdir(exist_ok=True, parents=True)
+    print(f"Copying examples into {config.data_store_path}")
+    _copy_examples(EXAMPLE_RULES, config.rules_store_path)
+    _copy_examples(EXAMPLE_GRAPHS, config.source_graph_path)
+    _copy_examples(EXAMPLE_WORKFLOWS, config.workflows_store_path)
+    config.staging_path.mkdir(exist_ok=True, parents=True)
 
 
-def create_data_dir_structure(target_data_dir: Path):
+def create_data_dir_structure(config: Config) -> None:
     """
     Create the data directory structure in empty directory
 
     Args:
         target_data_dir : The target directory
+        suffix : The suffix to add to the directory names
 
     """
-
-    (target_data_dir / "rules").mkdir(exist_ok=True, parents=True)
-    (target_data_dir / "source-graphs").mkdir(exist_ok=True, parents=True)
-    (target_data_dir / "staging").mkdir(exist_ok=True, parents=True)
-    (target_data_dir / "workflows").mkdir(exist_ok=True, parents=True)
+    for path in (
+        config.rules_store_path,
+        config.source_graph_path,
+        config.staging_path,
+        config.workflows_store_path,
+    ):
+        path.mkdir(exist_ok=True, parents=True)
 
 
 def _copy_examples(source_dir: Path, target_dir: Path):
@@ -44,3 +174,8 @@ def _copy_examples(source_dir: Path, target_dir: Path):
         if not (target := target_dir / relative).exists():
             target.parent.mkdir(exist_ok=True, parents=True)
             shutil.copy2(current, target)
+
+
+def configure_logging(level: str = "DEBUG", log_format: str = LOG_FORMAT):
+    """Configure logging based on config."""
+    logging.basicConfig(format=log_format, level=logging.getLevelName(level), datefmt=LOG_DATE_FORMAT)

cognite/neat/constants.py

@@ -9,6 +9,7 @@ PACKAGE_DIRECTORY = Path(neat.__file__).parent
 
 EXAMPLE_RULES = PACKAGE_DIRECTORY / "legacy" / "rules" / "examples"
 EXAMPLE_GRAPHS = PACKAGE_DIRECTORY / "legacy" / "graph" / "examples"
+_OLD_WORKFLOWS = PACKAGE_DIRECTORY / "legacy" / "workflows" / "examples"
 EXAMPLE_WORKFLOWS = PACKAGE_DIRECTORY / "workflows" / "examples"
 
 DEFAULT_NAMESPACE = Namespace("http://purl.org/cognite/neat#")

cognite/neat/graph/exceptions.py

@@ -1,5 +1,4 @@
-"""This module contains the definition of validation errors and warnings raised during graph methods
-"""
+"""This module contains the definition of validation errors and warnings raised during graph methods"""
 
 from cognite.neat.constants import DEFAULT_DOCS_URL
 from cognite.neat.exceptions import NeatException

cognite/neat/legacy/graph/exceptions.py

@@ -1,5 +1,4 @@
-"""This module contains the definition of validation errors and warnings raised during graph methods
-"""
+"""This module contains the definition of validation errors and warnings raised during graph methods"""
 
 from cognite.neat.constants import DEFAULT_DOCS_URL
 from cognite.neat.exceptions import NeatException

cognite/neat/legacy/graph/extractors/_mock_graph_generator.py

@@ -1,5 +1,4 @@
-"""This module is used to generate mock graph data for purposes of testing of NEAT.
-"""
+"""This module is used to generate mock graph data for purposes of testing of NEAT."""
 
 import logging
 import random

cognite/neat/legacy/graph/loaders/_asset_loader.py

@@ -119,12 +119,10 @@ class AssetLoader(CogniteLoader[AssetResource]):
         self._loaded_labels: set[str] = set()
 
     @overload
-    def load(self, stop_on_exception: Literal[True]) -> Iterable[AssetResource]:
-        ...
+    def load(self, stop_on_exception: Literal[True]) -> Iterable[AssetResource]: ...
 
     @overload
-    def load(self, stop_on_exception: Literal[False] = False) -> Iterable[AssetResource | ErrorDetails]:
-        ...
+    def load(self, stop_on_exception: Literal[False] = False) -> Iterable[AssetResource | ErrorDetails]: ...
 
     def load(self, stop_on_exception: bool = False) -> Iterable[AssetResource | ErrorDetails]:
         if self.rules.metadata.namespace is None:
@@ -259,12 +257,10 @@ class AssetLoader(CogniteLoader[AssetResource]):
             )
 
     @overload
-    def load_assets(self, stop_on_exception: Literal[True]) -> Iterable[AssetWrite]:
-        ...
+    def load_assets(self, stop_on_exception: Literal[True]) -> Iterable[AssetWrite]: ...
 
     @overload
-    def load_assets(self, stop_on_exception: Literal[False] = False) -> Iterable[AssetWrite | ErrorDetails]:
-        ...
+    def load_assets(self, stop_on_exception: Literal[False] = False) -> Iterable[AssetWrite | ErrorDetails]: ...
 
     def load_assets(self, stop_on_exception: Literal[True, False] = False) -> Iterable[AssetWrite | ErrorDetails]:
         for asset_resource in self.load(stop_on_exception):
@@ -272,14 +268,12 @@ class AssetLoader(CogniteLoader[AssetResource]):
                 yield asset_resource
 
     @overload
-    def load_relationships(self, stop_on_exception: Literal[True]) -> Iterable[RelationshipWrite]:
-        ...
+    def load_relationships(self, stop_on_exception: Literal[True]) -> Iterable[RelationshipWrite]: ...
 
     @overload
     def load_relationships(
         self, stop_on_exception: Literal[False] = False
-    ) -> Iterable[RelationshipWrite | ErrorDetails]:
-        ...
+    ) -> Iterable[RelationshipWrite | ErrorDetails]: ...
 
     def load_relationships(
         self, stop_on_exception: Literal[True, False] = False
@@ -348,7 +342,8 @@ class AssetLoader(CogniteLoader[AssetResource]):
             A dictionary with property type as key and a list of values as value.
         """
         properties_value_tuples: list[tuple[str, str]] = [
-            remove_namespace(prop, value) for _, prop, value in properties_values  # type: ignore[misc]
+            remove_namespace(prop, value)  # type: ignore[misc]
+            for _, prop, value in properties_values
         ]
         # We add an identifier which will be used as fallback for external_id
         properties_value_tuples.append((self._identifier, remove_namespace(instance_id)))

cognite/neat/legacy/graph/loaders/_base.py

@@ -26,12 +26,10 @@ class BaseLoader(ABC, Generic[T_Output]):
         self.graph_store = graph_store
 
     @overload
-    def load(self, stop_on_exception: Literal[True]) -> Iterable[T_Output]:
-        ...
+    def load(self, stop_on_exception: Literal[True]) -> Iterable[T_Output]: ...
 
     @overload
-    def load(self, stop_on_exception: Literal[False] = False) -> Iterable[T_Output | ErrorDetails]:
-        ...
+    def load(self, stop_on_exception: Literal[False] = False) -> Iterable[T_Output | ErrorDetails]: ...
 
     @abstractmethod
     def load(self, stop_on_exception: bool = False) -> Iterable[T_Output | ErrorDetails]:

cognite/neat/legacy/graph/loaders/_exceptions.py

@@ -1,6 +1,4 @@
-"""This module contains the definition of validation errors and warnings raised by various extractors
-"""
-
+"""This module contains the definition of validation errors and warnings raised by various extractors"""
 
 from pydantic_core import ErrorDetails, PydanticCustomError
 

cognite/neat/legacy/graph/loaders/core/rdf_to_assets.py

@@ -772,8 +772,7 @@ def categorize_assets(
     partitions: int = 2,
     stop_on_exception: bool = False,
     meta_keys: NeatMetadataKeys | None = None,
-) -> dict:
-    ...
+) -> dict: ...
 
 
 @overload
@@ -785,8 +784,7 @@ def categorize_assets(
     partitions: int = 2,
     stop_on_exception: bool = False,
     meta_keys: NeatMetadataKeys | None = None,
-) -> tuple[dict, dict]:
-    ...
+) -> tuple[dict, dict]: ...
 
 
 def categorize_assets(
@@ -1000,13 +998,11 @@ AssetLike: TypeAlias = Asset | dict[str, Any]
 
 
 @overload
-def remove_non_existing_labels(client: CogniteClient, assets: Sequence[AssetLike]) -> Sequence[AssetLike]:
-    ...
+def remove_non_existing_labels(client: CogniteClient, assets: Sequence[AssetLike]) -> Sequence[AssetLike]: ...
 
 
 @overload
-def remove_non_existing_labels(client: CogniteClient, assets: Mapping[str, AssetLike]) -> Mapping[str, AssetLike]:
-    ...
+def remove_non_existing_labels(client: CogniteClient, assets: Mapping[str, AssetLike]) -> Mapping[str, AssetLike]: ...
 
 
 def remove_non_existing_labels(

cognite/neat/legacy/graph/loaders/core/rdf_to_relationships.py

@@ -343,8 +343,7 @@ def categorize_relationships(
     data_set_id: int,
     return_report: Literal[False] = False,
     partitions: int = 40,
-) -> dict[str, list[Relationship] | list[RelationshipUpdate]]:
-    ...
+) -> dict[str, list[Relationship] | list[RelationshipUpdate]]: ...
 
 
 @overload
@@ -354,8 +353,7 @@ def categorize_relationships(
     data_set_id: int,
     return_report: Literal[True],
     partitions: int = 40,
-) -> tuple[dict[str, list[Relationship] | list[RelationshipUpdate]], dict[str, set]]:
-    ...
+) -> tuple[dict[str, list[Relationship] | list[RelationshipUpdate]], dict[str, set]]: ...
 
 
 def categorize_relationships(

cognite/neat/legacy/graph/loaders/rdf_to_dms.py

@@ -33,12 +33,10 @@ class DMSLoader(CogniteLoader[InstanceApply]):
         self.add_class_prefix = add_class_prefix
 
     @overload
-    def load(self, stop_on_exception: Literal[True]) -> Iterable[InstanceApply]:
-        ...
+    def load(self, stop_on_exception: Literal[True]) -> Iterable[InstanceApply]: ...
 
     @overload
-    def load(self, stop_on_exception: Literal[False] = False) -> Iterable[InstanceApply | ErrorDetails]:
-        ...
+    def load(self, stop_on_exception: Literal[False] = False) -> Iterable[InstanceApply | ErrorDetails]: ...
 
     def load(self, stop_on_exception: bool = False) -> Iterable[InstanceApply | ErrorDetails]:
         """Load the graph with data."""

cognite/neat/legacy/graph/loaders/validator.py

@@ -46,7 +46,7 @@ def _find_circular_reference_path(
 
 
 def validate_asset_hierarchy(
-    assets: dict[str, dict[str, Any]]
+    assets: dict[str, dict[str, Any]],
 ) -> tuple[list[str], list[list[str]], dict[str, list[str]]]:
     """Validates asset hierarchy and reports on orphan assets and circular dependency
 

cognite/neat/legacy/graph/transformations/transformer.py

@@ -1,5 +1,4 @@
-"""Methods to transform Domain Knowledge Graph to App Knowledge Graph
-"""
+"""Methods to transform Domain Knowledge Graph to App Knowledge Graph"""
 
 import logging
 import time

cognite/neat/legacy/rules/exporters/_rules2dms.py

@@ -1,5 +1,4 @@
-"""Exports rules to CDF Data Model Storage (DMS) through cognite-sdk.
-"""
+"""Exports rules to CDF Data Model Storage (DMS) through cognite-sdk."""
 
 import dataclasses
 import sys

cognite/neat/legacy/rules/exporters/_validation.py

@@ -15,13 +15,11 @@ from cognite.neat.legacy.rules.models.rules import (
 @overload
 def are_entity_names_dms_compliant(
     transformation_rules: Rules, return_report: Literal[True]
-) -> tuple[bool, list[dict]]:
-    ...
+) -> tuple[bool, list[dict]]: ...
 
 
 @overload
-def are_entity_names_dms_compliant(transformation_rules: Rules, return_report: Literal[False] = False) -> bool:
-    ...
+def are_entity_names_dms_compliant(transformation_rules: Rules, return_report: Literal[False] = False) -> bool: ...
 
 
 def are_entity_names_dms_compliant(
@@ -107,13 +105,11 @@ def are_entity_names_dms_compliant(
 
 
 @overload
-def are_properties_redefined(transformation_rules: Rules, return_report: Literal[True]) -> tuple[bool, list[dict]]:
-    ...
+def are_properties_redefined(transformation_rules: Rules, return_report: Literal[True]) -> tuple[bool, list[dict]]: ...
 
 
 @overload
-def are_properties_redefined(transformation_rules: Rules, return_report: Literal[False] = False) -> bool:
-    ...
+def are_properties_redefined(transformation_rules: Rules, return_report: Literal[False] = False) -> bool: ...
 
 
 def are_properties_redefined(

cognite/neat/legacy/rules/importers/_base.py

@@ -14,10 +14,6 @@ class BaseImporter(ABC):
     BaseImporter class which all importers inherit from.
     """
 
-    @abstractmethod
-    def __init__(self):
-        ...
-
     @abstractmethod
     def to_tables(self) -> dict[str, pd.DataFrame]:
         """Creates raw tables from the data."""

cognite/neat/legacy/rules/importers/_dms2rules.py

@@ -35,8 +35,6 @@ class DMSImporter(BaseImporter):
     """
 
     def __init__(self, views: Sequence[View] | DataModel[View], metadata: dict[str, str | float] | None = None):
-        super().__init__()
-
         if isinstance(views, DataModel):
             self.views = views.views
         else:

cognite/neat/legacy/rules/models/rdfpath.py

@@ -1,5 +1,4 @@
-"""
-"""
+""" """
 
 import re
 import sys

cognite/neat/legacy/workflows/examples/Export_DMS/workflow.yaml

@@ -0,0 +1,89 @@
+configs: []
+description: null
+implementation_module: null
+name: Export DMS
+steps:
+-   complex_configs: {}
+    configs: {}
+    description: null
+    enabled: true
+    id: step_861205
+    label: Upload Rules Spreadsheet
+    max_retries: 0
+    method: null
+    params:
+        file_type: rules
+    retry_delay: 3
+    stype: file_uploader
+    system_component_id: null
+    transition_to:
+    - step_295479
+    trigger: true
+    ui_config:
+        pos_x: 629
+        pos_y: 57
+-   complex_configs: {}
+    configs:
+        File name: ''
+        Report formatter: BasicHTML
+        Role: infer
+    description: null
+    enabled: true
+    id: step_295479
+    label: Validate
+    max_retries: 0
+    method: ExcelToRules
+    params: {}
+    retry_delay: 3
+    stype: stdstep
+    system_component_id: null
+    transition_to:
+    - step_50885
+    trigger: false
+    ui_config:
+        pos_x: 629
+        pos_y: 161
+-   complex_configs:
+        Components:
+            containers: true
+            data_models: true
+            spaces: true
+            views: true
+    configs:
+        Dry run: 'False'
+        Existing component handling: update
+        Multi-space components create: 'True'
+    description: null
+    enabled: true
+    id: step_50885
+    label: Export Data Model to CDF
+    max_retries: 0
+    method: RulesToDMS
+    params: {}
+    retry_delay: 3
+    stype: stdstep
+    system_component_id: null
+    transition_to:
+    - step_171560
+    trigger: false
+    ui_config:
+        pos_x: 629
+        pos_y: 243
+-   complex_configs: {}
+    configs: {}
+    description: null
+    enabled: true
+    id: step_171560
+    label: Export Transformations
+    max_retries: 0
+    method: RulesToCDFTransformations
+    params: {}
+    retry_delay: 3
+    stype: stdstep
+    system_component_id: null
+    transition_to: []
+    trigger: false
+    ui_config:
+        pos_x: 629
+        pos_y: 342
+system_components: []