diff --git a/cognite/client/_api/data_modeling/graphql.py b/cognite/client/_api/data_modeling/graphql.py index 3f2c9a0e16..52dec1c4d3 100644 --- a/cognite/client/_api/data_modeling/graphql.py +++ b/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"]) diff --git a/tests/tests_integration/test_api/test_data_modeling/test_graphql.py b/tests/tests_integration/test_api/test_data_modeling/test_graphql.py index c9c6fa2024..0e0f6fef0b 100644 --- a/tests/tests_integration/test_api/test_data_modeling/test_graphql.py +++ b/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()