Add hidden method for wiping and regenerating DML (#1420)

cognite/client/_api/data_modeling/graphql.py

@@ -11,16 +11,49 @@
 
 
 class DataModelingGraphQLAPI(APIClient):
-    def _post_graphql(self, url_path: str, json: dict) -> dict[str, Any]:
+    def _post_graphql(self, url_path: str, query_name: str, json: dict) -> dict[str, Any]:
         res = self._post(url_path=url_path, json=json).json()
         # Errors can be passed both at top level and nested in the response:
-        errors = res.get("errors", []) + (
-            (res.get("data", {}).get("upsertGraphQlDmlVersion") or {}).get("errors") or []
-        )
+        errors = res.get("errors", []) + ((res.get("data", {}).get(query_name) or {}).get("errors") or [])
         if errors:
             raise CogniteGraphQLError([GraphQLErrorSpec.load(error) for error in errors])
         return res["data"]
 
+    def _unsafely_wipe_and_regenerate_dml(self, id: DataModelIdentifier) -> str:
+        """Wipe and regenerate the DML for a given data model.
+
+        Note:
+            This removes all comments from the DML.
+
+        Args:
+            id (DataModelIdentifier): The data model to apply DML to.
+
+        Returns:
+            str: The new DML
+        """
+        graphql_body = """
+            query WipeAndRegenerateDml($space: String!, $externalId: String!, $version: String!) {
+                unsafelyWipeAndRegenerateDmlBasedOnDataModel(space: $space, externalId: $externalId, version: $version) {
+                    items {
+                        graphQlDml
+                    }
+                }
+            }
+        """
+        data_model_id = DataModelId.load(id)
+        payload = {
+            "query": textwrap.dedent(graphql_body),
+            "variables": {
+                "space": data_model_id.space,
+                "externalId": data_model_id.external_id,
+                "version": data_model_id.version,
+            },
+        }
+
+        query_name = "unsafelyWipeAndRegenerateDmlBasedOnDataModel"
+        res = self._post_graphql(url_path="/dml/graphql", query_name=query_name, json=payload)
+        return res[query_name]["items"][0]["graphQlDml"]
+
     def apply_dml(
         self,
         id: DataModelIdentifier,
@@ -96,5 +129,7 @@ def apply_dml(
                 }
             },
         }
-        res = self._post_graphql(url_path="/dml/graphql", json=payload)
-        return DMLApplyResult.load(res["upsertGraphQlDmlVersion"]["result"])
+
+        query_name = "upsertGraphQlDmlVersion"
+        res = self._post_graphql(url_path="/dml/graphql", query_name=query_name, json=payload)
+        return DMLApplyResult.load(res[query_name]["result"])

tests/tests_integration/test_api/test_data_modeling/test_graphql.py

@@ -1,3 +1,5 @@
+import textwrap
+
 import pytest
 
 from cognite.client import CogniteClient
@@ -49,3 +51,16 @@ def test_apply_dm_raise_top_level_error(self, cognite_client: CogniteClient, int
         exception_message = str(exc.value)
         assert "version" in exception_message
         assert "invalid value" in exception_message
+
+    def test_wipe_and_regenerate_dml(self, cognite_client: CogniteClient, data_model: DataModel) -> None:
+        res = cognite_client.data_modeling.graphql._unsafely_wipe_and_regenerate_dml(data_model.as_id())
+        expected = """
+            type SomeType @view(version: "31e24bb338352b") {
+              someProp: String!
+            }
+
+            type AnotherType @view(version: "855f080c9de22f") {
+              anotherProp: String!
+            }
+        """
+        assert res.strip() == textwrap.dedent(expected).strip()