[FEATURE] Update expectations and checkpoints v1 stores to implement gx_cloud_response_json_to_object_collection

great_expectations/data_context/store/checkpoint_store.py

@@ -247,8 +247,8 @@ def get_key(self, name: str, id: str | None = None) -> GXCloudIdentifier | Strin
         return self._key_class(key=name)
 
     @override
-    @staticmethod
-    def gx_cloud_response_json_to_object_dict(response_json: dict) -> dict:
+    @classmethod
+    def gx_cloud_response_json_to_object_dict(cls, response_json: dict) -> dict:
         response_data = response_json["data"]
 
         checkpoint_data: dict
@@ -263,8 +263,13 @@ def gx_cloud_response_json_to_object_dict(response_json: dict) -> dict:
         else:
             checkpoint_data = response_data
 
-        id: str = checkpoint_data["id"]
-        checkpoint_config_dict: dict = checkpoint_data["attributes"]["checkpoint_config"]
+        return cls._convert_raw_json_to_object_dict(checkpoint_data)
+
+    @override
+    @staticmethod
+    def _convert_raw_json_to_object_dict(data: dict) -> dict:
+        id: str = data["id"]
+        checkpoint_config_dict: dict = data["attributes"]["checkpoint_config"]
         checkpoint_config_dict["id"] = id
 
         return checkpoint_config_dict

great_expectations/data_context/store/datasource_store.py

@@ -138,8 +138,9 @@ def deserialize(
             return self._schema.loads(value)
 
     @override
-    @staticmethod
+    @classmethod
     def gx_cloud_response_json_to_object_dict(
+        cls,
         response_json: CloudResponsePayloadTD,  # type: ignore[override]
     ) -> dict:
         """
@@ -158,22 +159,7 @@ def gx_cloud_response_json_to_object_dict(
 
     @override
     @staticmethod
-    def gx_cloud_response_json_to_object_collection(
-        response_json: CloudResponsePayloadTD,  # type: ignore[override]
-    ) -> list[dict]:
-        """
-        This method takes full json response from GX cloud and outputs a list of dicts appropriate for
-        deserialization into a collection of GX objects
-        """  # noqa: E501
-        logger.debug(f"GE Cloud Response JSON ->\n{pf(response_json, depth=3)}")
-        data = response_json["data"]
-        if not isinstance(data, list):
-            raise TypeError("GX Cloud did not return a collection of Datasources when expected")  # noqa: TRY003
-
-        return [DatasourceStore._convert_raw_json_to_object_dict(d) for d in data]
-
-    @staticmethod
-    def _convert_raw_json_to_object_dict(data: DataPayload) -> dict:
+    def _convert_raw_json_to_object_dict(data: DataPayload) -> dict:  # type: ignore[override]
         datasource_id: str = data["id"]
         datasource_config_dict: dict = data["attributes"]["datasource_config"]
         datasource_config_dict["id"] = datasource_id

great_expectations/data_context/store/expectations_store.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import uuid
-from typing import TYPE_CHECKING, Dict, List, Optional, TypeVar, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, TypeVar, Union
 
 import great_expectations.exceptions as gx_exceptions
 from great_expectations.compatibility import pydantic
@@ -111,8 +111,8 @@ def __init__(
         filter_properties_dict(properties=self._config, clean_falsy=True, inplace=True)
 
     @override
-    @staticmethod
-    def gx_cloud_response_json_to_object_dict(response_json: dict) -> dict:
+    @classmethod
+    def gx_cloud_response_json_to_object_dict(cls, response_json: dict) -> dict:
         """
         This method takes full json response from GX cloud and outputs a dict appropriate for
         deserialization into a GX object
@@ -129,13 +129,18 @@ def gx_cloud_response_json_to_object_dict(response_json: dict) -> dict:
         else:
             suite_data = response_json["data"]
 
+        return cls._convert_raw_json_to_object_dict(suite_data)
+
+    @override
+    @staticmethod
+    def _convert_raw_json_to_object_dict(data: dict[str, Any]) -> dict[str, Any]:
         # Cloud backend adds a default result format type of None, so ensure we remove it:
-        for expectation in suite_data.get("expectations", []):
+        for expectation in data.get("expectations", []):
             kwargs = expectation["kwargs"]
             if "result_format" in kwargs and kwargs["result_format"] is None:
                 kwargs.pop("result_format")
 
-        suite_dto = ExpectationSuiteDTO.parse_obj(suite_data)
+        suite_dto = ExpectationSuiteDTO.parse_obj(data)
         return suite_dto.dict()
 
     def add_expectation(self, suite: ExpectationSuite, expectation: _TExpectation) -> _TExpectation:

great_expectations/data_context/store/store.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+from pprint import pformat as pf
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -102,20 +103,38 @@
             )
         self._use_fixed_length_key = self._store_backend.fixed_length_key
 
-    @staticmethod
-    def gx_cloud_response_json_to_object_dict(response_json: Dict) -> Dict:
+    @classmethod
+    def gx_cloud_response_json_to_object_dict(cls, response_json: Dict) -> Dict:
         """
         This method takes full json response from GX cloud and outputs a dict appropriate for
         deserialization into a GX object
         """
         return response_json
 
-    @staticmethod
-    def gx_cloud_response_json_to_object_collection(response_json: Dict) -> List[Dict]:
+    @classmethod
+    def gx_cloud_response_json_to_object_collection(cls, response_json: Dict) -> List[Dict]:
         """
         This method takes full json response from GX cloud and outputs a list of dicts appropriate for
         deserialization into a collection of GX objects
         """  # noqa: E501
+        logger.debug(f"GE Cloud Response JSON ->\n{pf(response_json, depth=3)}")
+        data = response_json["data"]
+        if not isinstance(data, list):
+            raise TypeError("GX Cloud did not return a collection of Datasources when expected")  # noqa: TRY003
+
+        return [cls._convert_raw_json_to_object_dict(d) for d in data]
+
+    @staticmethod
+    def _convert_raw_json_to_object_dict(data: dict[str, Any]) -> dict[str, Any]:
+        """Method to convert data from API to raw object dict
+
+        This SHOULD be used by both gx_cloud_response_json_to_object_collection
+        and gx_cloud_response_json_to_object_dict. It is a means of keeping
+        response parsing DRY for different response types, e.g. collections
+        may be shaped like {"data": [item1, item2, ...]} while single items
+        may be shaped like {"data": item}. This allows for pulling out the
+        data key and passing it to the appropriate method for conversion.
+        """
         raise NotImplementedError
 
     def _validate_key(self, key: DataContextKey) -> None:

tests/data_context/store/test_datasource_store.py

@@ -646,7 +646,7 @@ def test_gx_cloud_response_json_to_object_dict(
         assert actual == expected
 
 
-@pytest.mark.cloud
+@pytest.mark.unit
 def test_gx_cloud_response_json_to_object_collection():
     response_json = {
         "data": [

tests/data_context/store/test_expectations_store.py

@@ -179,86 +179,53 @@ def test_expectations_store_report_same_id_with_same_configuration_TupleFilesyst
     assert gen_directory_tree_str(project_path) == initialized_directory_tree_with_store_backend_id
 
 
-@pytest.mark.cloud
+def _create_suite_config(name: str, id: str, expectations: list[dict] | None = None) -> dict:
+    return {
+        "id": id,
+        "name": name,
+        "expectations": expectations or [],
+        "meta": {},
+        "notes": None,
+    }
+
+
+_SUITE_CONFIG = _create_suite_config("my_suite", "03d61d4e-003f-48e7-a3b2-f9f842384da3")
+_SUITE_CONFIG_WITH_EXPECTATIONS = _create_suite_config(
+    "my_suite_with_expectations",
+    "03d61d4e-003f-48e7-a3b2-f9f842384da3",
+    [
+        {
+            "expectation_type": "expect_column_to_exist",
+            "id": "c8a239a6-fb80-4f51-a90e-40c38dffdf91",
+            "kwargs": {"column": "infinities"},
+            "meta": {},
+            "expectation_context": None,
+            "rendered_content": [],
+        }
+    ],
+)
+
+
+@pytest.mark.unit
 @pytest.mark.parametrize(
     "response_json, expected, error_type",
     [
         pytest.param(
-            {
-                "data": {
-                    "id": "03d61d4e-003f-48e7-a3b2-f9f842384da3",
-                    "name": "my_suite",
-                    "expectations": [],
-                    "meta": {},
-                    "notes": None,
-                }
-            },
-            {
-                "name": "my_suite",
-                "id": "03d61d4e-003f-48e7-a3b2-f9f842384da3",
-                "expectations": [],
-                "meta": {},
-                "notes": None,
-            },
+            {"data": _SUITE_CONFIG},
+            _SUITE_CONFIG,
             None,
             id="single_config",
         ),
         pytest.param({"data": []}, None, ValueError, id="empty_payload"),
         pytest.param(
-            {
-                "data": [
-                    {
-                        "id": "03d61d4e-003f-48e7-a3b2-f9f842384da3",
-                        "name": "my_suite",
-                        "expectations": [],
-                        "meta": {},
-                        "notes": None,
-                    }
-                ]
-            },
-            {
-                "id": "03d61d4e-003f-48e7-a3b2-f9f842384da3",
-                "name": "my_suite",
-                "expectations": [],
-                "meta": {},
-                "notes": None,
-            },
+            {"data": [_SUITE_CONFIG]},
+            _SUITE_CONFIG,
             None,
             id="single_config_in_list",
         ),
         pytest.param(
-            {
-                "data": {
-                    "id": "03d61d4e-003f-48e7-a3b2-f9f842384da3",
-                    "name": "my_suite",
-                    "expectations": [
-                        {
-                            "expectation_type": "expect_column_to_exist",
-                            "id": "c8a239a6-fb80-4f51-a90e-40c38dffdf91",
-                            "kwargs": {"column": "infinities", "result_format": None},
-                            "meta": {},
-                        }
-                    ],
-                    "meta": {},
-                    "notes": None,
-                }
-            },
-            {
-                "name": "my_suite",
-                "id": "03d61d4e-003f-48e7-a3b2-f9f842384da3",
-                "expectations": [
-                    {
-                        "expectation_type": "expect_column_to_exist",
-                        "id": "c8a239a6-fb80-4f51-a90e-40c38dffdf91",
-                        "kwargs": {"column": "infinities"},
-                        "meta": {},
-                        "expectation_context": None,
-                        "rendered_content": [],
-                    }
-                ],
-                "meta": {},
-                "notes": None,
-            },
+            {"data": _SUITE_CONFIG_WITH_EXPECTATIONS},
+            _SUITE_CONFIG_WITH_EXPECTATIONS,
             None,
             id="null_result_format",
         ),
@@ -275,6 +242,21 @@ def test_gx_cloud_response_json_to_object_dict(
         assert actual == expected
 
 
+@pytest.mark.unit
+def test_gx_cloud_response_json_to_object_collection():
+    response_json = {
+        "data": [
+            _SUITE_CONFIG,
+            _SUITE_CONFIG_WITH_EXPECTATIONS,
+        ]
+    }
+
+    result = ExpectationsStore.gx_cloud_response_json_to_object_collection(response_json)
+
+    expected = [_SUITE_CONFIG, _SUITE_CONFIG_WITH_EXPECTATIONS]
+    assert result == expected
+
+
 @pytest.mark.unit
 def test_get_key_in_non_cloud_mode(empty_data_context):
     name = "test-name"