Merge pull request #437 from cognitedata/main

Release 1.0.2

cognite/pygen/_core/templates/_api_client_multi_model.py.jinja

@@ -28,6 +28,7 @@ class {{ api.name }}:
         version: {{ api.model.version }}
 
     """
+    _data_model_id = dm.DataModelId("{{ data_model.space }}", "{{ data_model.external_id }}", "{{ data_model.version }}")
 
     def __init__(self, client: CogniteClient):
         self._client = client

cognite/pygen/_core/templates/_api_client_single_model.py.jinja

@@ -36,6 +36,7 @@ class {{ client_name }}:
         externalId: {{ data_model.external_id }}
         version: {{ data_model.version }}
     """
+    _data_model_id = dm.DataModelId("{{ data_model.space }}", "{{ data_model.external_id }}", "{{ data_model.version }}")
 
     def __init__(self, config_or_client: CogniteClient | ClientConfig):
         if isinstance(config_or_client, CogniteClient):

cognite/pygen/_core/templates/config.py.jinja

@@ -1,4 +1,4 @@
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 
 class GlobalConfig(BaseModel, validate_assignment=True):
@@ -8,10 +8,12 @@ class GlobalConfig(BaseModel, validate_assignment=True):
         validate_retrieve (bool): Whether to validate the data retrieved from the CDF. Defaults to True.
             Note setting this to False can lead to unexpected behavior if required fields are missing
             or have the wrong type.
+        max_select_depth (int): The maximum depth of select queries. Defaults to 4.
 
     """
 
     validate_retrieve: bool = True
+    max_select_depth: int = Field(3, ge=1)
 
 
 global_config = GlobalConfig()

cognite/pygen/_core/templates/data_class_node.py.jinja

@@ -34,6 +34,7 @@ from cognite.client.data_classes import (
 {% endif %}
 from pydantic import field_validator, model_validator, ValidationInfo
 
+from {{ top_level_package }}.config import global_config
 from {{ top_level_package }}.data_classes._core import (
     {% if has_default_instance_space %}
     DEFAULT_INSTANCE_SPACE,
@@ -387,7 +388,7 @@ class _{{ data_class.query_cls_name }}(NodeQueryCore[T_DomainModelList, {{ data_
 
         {% for field in data_class.fields_of_type(ft.BaseConnectionField) %}
         {% if not field.is_direct_relation_no_source %}
-        if _{{ field.linked_class.query_cls_name }} not in created_types:
+        if _{{ field.linked_class.query_cls_name }} not in created_types and len(creation_path) + 1 < global_config.max_select_depth:
             self.{{ field.name }} = _{{ field.linked_class.query_cls_name }}(
                 created_types.copy(),
                 self._creation_path,

cognite/pygen/_core/templates/data_classes_core_query_select.py.jinja

@@ -15,6 +15,7 @@ from typing import (
 from cognite.client import CogniteClient
 from cognite.client import data_modeling as dm
 
+from {{top_level_package}}.config import global_config
 from {{top_level_package}}.data_classes._core.query.filter_classes import Filtering
 from {{top_level_package}}.data_classes._core.base import (
     DomainModelList,
@@ -78,6 +79,17 @@ class QueryCore(Generic[T_DomainList, T_DomainListEnd]):
             raise ValueError(f"Circular reference detected. Cannot query a circular reference: {nodes}")
         elif self._connection_type == "reverse-list":
             raise ValueError(f"Cannot query across a reverse-list connection.")
+        elif len(self._creation_path) >= global_config.max_select_depth:
+            hint = f"""You can increase the max_select_depth in the global config.
+```
+from omni.config import global_config
+
+global_config.max_select_depth = {global_config.max_select_depth+1}
+```
+"""
+            raise ValueError(
+                f"Max select depth reached. Cannot query deeper than {global_config.max_select_depth}.\n{hint}"
+            )
         error_message = f"'{self.__class__.__name__}' object has no attribute '{item}'"
         attributes = [name for name in vars(self).keys() if not name.startswith("_")]
         if matches := difflib.get_close_matches(item, attributes):

cognite/pygen/_generator.py

@@ -163,6 +163,7 @@ def _generate_sdk(
     if return_sdk_files:
         return sdk
     output_dir = output_dir or (Path.cwd() / Path(top_level_package.replace(".", "/")))
+    output_dir = Path(output_dir)
     logger(f"Writing SDK to {output_dir}")
     write_sdk_to_disk(sdk, output_dir, overwrite, logger, format_code)
     logger("Done!")

cognite/pygen/_version.py

@@ -1 +1 @@
-__version__ = "1.0.1"
+__version__ = "1.0.2"

docs/CHANGELOG.md

@@ -13,6 +13,13 @@ Changes are grouped as follows
 - `Fixed` for any bug fixes.
 - `Security` in case of vulnerabilities.
 
+## [1.0.2] - 2025-02-24
+### Fixed
+- Calling `.select()` on large, highly connected data models no longer raises a `MemoryExceptionError`. This was caused
+  by the underlying data structure created by pygen increasing in size exponentially with the number of connections.
+  This is now fixed by introducing a `global_config.max_select_depth` that limits the size of the data structure
+  pygen created.
+
 ## [1.0.1] - 2025-02-13
 ### Fixed
 - In a Pyodide environment with a large model, the `generate_sdk_notebook` can hit a recursion limit in pydantic