Repository: great-expectations/great_expectations
Pull request #6962 (2023-01-27 15:21:04)

[MAINTENANCE] Standardize `Checkpoint` CRUD
Changes proposed in this pull request:
- Add `update_checkpoint` and `add_or_update_checkpoint`
- Misc cleanup


### Definition of Done
- [x] My code follows the Great Expectations [style guide](https://docs.greatexpectations.io/docs/contributing/style_guides/code_style)
- [x] I have performed a [self-review](https://docs.greatexpectations.io/docs/contributing/contributing_checklist) of my own code
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] I have added [unit tests](https://docs.greatexpectations.io/docs/contributing/contributing_test#writing-unit-and-integration-tests) where applicable and made sure that new and existing tests are passing.
- [x] I have run any local integration tests and made sure that nothing is broken.


codesee-maps[bot]: <details>
<summary><b>👇 <a href="https://app.codesee.io/r/reviews?pr=6962&src=https%3A%2F%2Fgithub.com%2Fgreat-expectations%2Fgreat_expectations">Click on the image</a> for a new way to code review</b>
</summary>

  - Make big changes easier — review code in small groups of related files
  - Know where to start — see the whole change at a glance
  - Take a code tour — explore the change with an interactive tour
  - Make comments and review — all fully sync’ed with github

    [Try it now!](https://app.codesee.io/r/reviews?pr=6962&src=https%3A%2F%2Fgithub.com%2Fgreat-expectations%2Fgreat_expectations)
</details>

[![Review these changes using an interactive CodeSee Map](https://s3.us-east-2.amazonaws.com/maps.codesee.io/images/github/great-expectations/great_expectations/6962/83d61197/0accfac7182639dc5cd04ccde1b86d0b1053b659.svg)](https://app.codesee.io/r/reviews?pr=6962&src=https%3A%2F%2Fgithub.com%2Fgreat-expectations%2Fgreat_expectations)

#### Legend
![CodeSee Map Legend](https://app.codesee.io/Legend_v2.png)
netlify[bot]: ### <span aria-hidden="true">👷</span> Deploy Preview for *niobium-lead-7998* processing.


|  Name | Link |
|---------------------------------|------------------------|
|<span aria-hidden="true">🔨</span> Latest commit | 0accfac7182639dc5cd04ccde1b86d0b1053b659 |
|<span aria-hidden="true">🔍</span> Latest deploy log | https://app.netlify.com/sites/niobium-lead-7998/deploys/63d3ebf7eca5890008c5840c |

Patch:
diff --git a/great_expectations/data_context/data_context/abstract_data_context.py b/great_expectations/data_context/data_context/abstract_data_context.py
index 19580a8f9..bba3998fc 100644
--- a/great_expectations/data_context/data_context/abstract_data_context.py
+++ b/great_expectations/data_context/data_context/abstract_data_context.py
@@ -1530,8 +1530,98 @@ class AbstractDataContext(ConfigPeer, ABC):
             default_validation_id=default_validation_id,
         )
 
-        self.checkpoint_store.add_checkpoint(checkpoint=checkpoint, name=name, id=id)
-        return checkpoint
+        return self.checkpoint_store.add_checkpoint(checkpoint)
+
+    def update_checkpoint(self, checkpoint: Checkpoint) -> Checkpoint:
+        """Update a Checkpoint that already exists.
+
+        Args:
+            checkpoint: The checkpoint to use to update.
+
+        Raises:
+            DataContextError: A suite with the given name does not already exist.
+        """
+        return self.checkpoint_store.update_checkpoint(checkpoint)
+
+    def add_or_update_checkpoint(  # noqa: C901 - Complexity 23
+        self,
+        name: str,
+        id: str | None = None,
+        config_version: int | float | None = None,
+        template_name: str | None = None,
+        module_name: str | None = None,
+        class_name: str | None = None,
+        run_name_template: str | None = None,
+        expectation_suite_name: str | None = None,
+        batch_request: dict | None = None,
+        action_list: list[dict] | None = None,
+        evaluation_parameters: dict | None = None,
+        runtime_configuration: dict | None = None,
+        validations: list[dict] | None = None,
+        profilers: list[dict] | None = None,
+        # the following four arguments are used by SimpleCheckpoint
+        site_names: str | list[str] | None = None,
+        slack_webhook: str | None = None,
+        notify_on: str | None = None,
+        notify_with: str | list[str] | None = None,
+        expectation_suite_id: str | None = None,
+        default_validation_id: str | None = None,
+    ) -> Checkpoint:
+        """Add a new Checkpoint or update an existing one on the context depending on whether it already exists or not.
+
+        Args:
+            name: The name to give the checkpoint.
+            id: The ID to associate with this checkpoint.
+            config_version: The config version of this checkpoint.
+            template_name: The template to use in generating this checkpoint.
+            module_name: The module name to use in generating this checkpoint.
+            class_name: The class name to use in generating this checkpoint.
+            run_name_template: The run name template to use in generating this checkpoint.
+            expectation_suite_name: The expectation suite name to use in generating this checkpoint.
+            batch_request: The batch request to use in generating this checkpoint.
+            action_list: The action list to use in generating this checkpoint.
+            evaluation_parameters: The evaluation parameters to use in generating this checkpoint.
+            runtime_configuration: The runtime configuration to use in generating this checkpoint.
+            validations: The validations to use in generating this checkpoint.
+            profilers: The profilers to use in generating this checkpoint.
+            site_names: The site names to use in generating this checkpoint. This is only used for SimpleCheckpoint configuration.
+            slack_webhook: The slack webhook to use in generating this checkpoint. This is only used for SimpleCheckpoint configuration.
+            notify_on: The notify on setting to use in generating this checkpoint. This is only used for SimpleCheckpoint configuration.
+            notify_with: The notify with setting to use in generating this checkpoint. This is only used for SimpleCheckpoint configuration.
+            expectation_suite_id: The expectation suite GE Cloud ID to use in generating this checkpoint.
+            default_validation_id: The default validation ID to use in generating this checkpoint.
+
+        Returns:
+            A new Checkpoint or an updated once (depending on whether or not it existed before this method call).
+        """
+        from great_expectations.checkpoint.checkpoint import Checkpoint
+
+        checkpoint = Checkpoint.construct_from_config_args(
+            data_context=self,
+            checkpoint_store_name=self.checkpoint_store_name,  # type: ignore[arg-type]
+            name=name,
+            config_version=config_version,
+            template_name=template_name,
+            module_name=module_name,
+            class_name=class_name,
+            run_name_template=run_name_template,
+            expectation_suite_name=expectation_suite_name,
+            batch_request=batch_request,
+            action_list=action_list,
+            evaluation_parameters=evaluation_parameters,
+            runtime_configuration=runtime_configuration,
+            validations=validations,
+            profilers=profilers,
+            # the following four arguments are used by SimpleCheckpoint
+            site_names=site_names,
+            slack_webhook=slack_webhook,
+            notify_on=notify_on,
+            notify_with=notify_with,
+            ge_cloud_id=id,
+            expectation_suite_ge_cloud_id=expectation_suite_id,
+            default_validation_id=default_validation_id,
+        )
+        return self.checkpoint_store.add_or_update_checkpoint(checkpoint)
 
     def get_checkpoint(
         self,
diff --git a/great_expectations/data_context/store/checkpoint_store.py b/great_expectations/data_context/store/checkpoint_store.py
index dc6ff5256..3eabd4146 100644
--- a/great_expectations/data_context/store/checkpoint_store.py
+++ b/great_expectations/data_context/store/checkpoint_store.py
@@ -168,17 +168,44 @@ class CheckpointStore(ConfigurationStore):
 
         return checkpoint_config
 
-    def add_checkpoint(
-        self, checkpoint: Checkpoint, name: str | None, id: str | None
-    ) -> None:
-        key: Union[GXCloudIdentifier, ConfigurationIdentifier] = self.determine_key(
-            name=name, ge_cloud_id=id
-        )
-        checkpoint_config: CheckpointConfig = checkpoint.get_config()  # type: ignore[assignment]
+    def add_checkpoint(self, checkpoint: Checkpoint) -> Checkpoint:
+        key = self._construct_key_from_checkpoint(checkpoint)
+        return self._add_checkpoint(key=key, checkpoint=checkpoint)
+
+    def update_checkpoint(self, checkpoint: Checkpoint) -> Checkpoint:
+        key = self._construct_key_from_checkpoint(checkpoint)
+        if not self.has_key(key):
+            raise gx_exceptions.CheckpointNotFoundError(
+                f"Could not find a Checkpoint named {checkpoint.name}"
+            )
+        return self._add_checkpoint(key=key, checkpoint=checkpoint)
+
+    def add_or_update_checkpoint(self, checkpoint: Checkpoint) -> Checkpoint:
+        key = self._construct_key_from_checkpoint(checkpoint)
+        if self.has_key(key):
+            return self.update_checkpoint(checkpoint)
+        return self._add_checkpoint(key=key, checkpoint=checkpoint)
+
+    def _construct_key_from_checkpoint(
+        self, checkpoint: Checkpoint
+    ) -> Union[GXCloudIdentifier, ConfigurationIdentifier]:
+        name = checkpoint.name
+        id = checkpoint.ge_cloud_id
+        if id:
+            return self.determine_key(ge_cloud_id=str(id))
+        return self.determine_key(name=name)
+
+    def _add_checkpoint(
+        self,
+        key: GXCloudIdentifier | ConfigurationIdentifier,
+        checkpoint: Checkpoint,
+    ) -> Checkpoint:
+        checkpoint_config = checkpoint.get_config()
         checkpoint_ref = self.set(key=key, value=checkpoint_config)  # type: ignore[func-returns-value]
         if isinstance(checkpoint_ref, GXCloudIDAwareRef):
             cloud_id = checkpoint_ref.cloud_id
             checkpoint.ge_cloud_id = uuid.UUID(cloud_id)  # type: ignore[misc]
+        return checkpoint
 
     def create(self, checkpoint_config: CheckpointConfig) -> Optional[DataContextKey]:
         """Create a checkpoint config in the store using a store_backend-specific key.
diff --git a/great_expectations/data_context/store/configuration_store.py b/great_expectations/data_context/store/configuration_store.py
index 0522eff95..3783121f7 100644
--- a/great_expectations/data_context/store/configuration_store.py
+++ b/great_expectations/data_context/store/configuration_store.py
@@ -151,7 +151,7 @@ class ConfigurationStore(Store):
         raise NotImplementedError
 
     def determine_key(
-        self, name: Optional[str], ge_cloud_id: Optional[str]
+        self, name: Optional[str] = None, ge_cloud_id: Optional[str] = None
     ) -> Union[GXCloudIdentifier, ConfigurationIdentifier]:
         assert bool(name) ^ bool(
             ge_cloud_id


Test patch:
diff --git a/tests/data_context/store/test_checkpoint_store.py b/tests/data_context/store/test_checkpoint_store.py
index 62c9c6896..f24f4f3a9 100644
--- a/tests/data_context/store/test_checkpoint_store.py
+++ b/tests/data_context/store/test_checkpoint_store.py
@@ -520,11 +520,12 @@ def test_add_checkpoint(
 
     context = mock.MagicMock(spec=FileDataContext)
     context._usage_statistics_handler = mock.MagicMock()
-    checkpoint = Checkpoint(name="my_checkpoint", data_context=context)
+    checkpoint_name = "my_checkpoint"
+    checkpoint = Checkpoint(name=checkpoint_name, data_context=context)
 
-    store.add_checkpoint(checkpoint=checkpoint, name="my_checkpoint", id=None)
+    store.add_checkpoint(checkpoint=checkpoint)
 
     mock_backend.set.assert_called_once_with(
-        ("my_checkpoint",),
+        (checkpoint_name,),
         "name: my_checkpoint\nconfig_version:\nmodule_name: great_expectations.checkpoint\nclass_name: LegacyCheckpoint\n",
     )
diff --git a/tests/data_context/test_data_context_state_management.py b/tests/data_context/test_data_context_state_management.py
index 466cb62a4..afedaf1f8 100644
--- a/tests/data_context/test_data_context_state_management.py
+++ b/tests/data_context/test_data_context_state_management.py
@@ -5,13 +5,16 @@ from typing import Mapping
 import pytest
 
 import great_expectations.exceptions as gx_exceptions
+from great_expectations.checkpoint.checkpoint import Checkpoint
 from great_expectations.core.expectation_configuration import ExpectationConfiguration
 from great_expectations.core.expectation_suite import ExpectationSuite
 from great_expectations.data_context.data_context.ephemeral_data_context import (
     EphemeralDataContext,
 )
 from great_expectations.data_context.store import ExpectationsStore, ProfilerStore
+from great_expectations.data_context.store.checkpoint_store import CheckpointStore
 from great_expectations.data_context.types.base import (
+    CheckpointConfig,
     DataContextConfig,
     DatasourceConfig,
     InMemoryStoreBackendDefaults,
@@ -44,6 +47,19 @@ class ProfilerStoreSpy(ProfilerStore):
         return super().set(key=key, value=value, **kwargs)
 
 
+class CheckpointStoreSpy(CheckpointStore):
+
+    STORE_NAME = "checkpoint_store"
+
+    def __init__(self) -> None:
+        self.save_count = 0
+        super().__init__(CheckpointStoreSpy.STORE_NAME)
+
+    def set(self, key, value, **kwargs):
+        self.save_count += 1
+        return super().set(key=key, value=value, **kwargs)
+
+
 class EphemeralDataContextSpy(EphemeralDataContext):
     """
     Simply wraps around EphemeralDataContext but keeps tabs on specific method calls around state management.
@@ -57,6 +73,7 @@ class EphemeralDataContextSpy(EphemeralDataContext):
         self.save_count = 0
         self._expectations_store = ExpectationsStoreSpy()
         self._profiler_store = ProfilerStoreSpy()
+        self._checkpoint_store = CheckpointStoreSpy()
 
     @property
     def expectations_store(self):
@@ -66,6 +83,10 @@ class EphemeralDataContextSpy(EphemeralDataContext):
     def profiler_store(self):
         return self._profiler_store
 
+    @property
+    def checkpoint_store(self):
+        return self._checkpoint_store
+
     def _save_project_config(self):
         """
         No-op our persistence mechanism but increment an internal counter to ensure it was used.
@@ -548,3 +569,189 @@ def test_add_or_update_profiler_updates_successfully(
     # Rules get converted to a list within the RBP constructor
     assert sorted(rule.name for rule in profiler.rules) == sorted(new_rules.keys())
     assert context.profiler_store.save_count == 2
+
+
+@pytest.mark.unit
+def test_update_checkpoint_success(
+    in_memory_data_context: EphemeralDataContextSpy,
+    checkpoint_config: dict,
+):
+    context = in_memory_data_context
+    name = "my_checkpoint"
+    checkpoint_config["name"] = name
+
+    checkpoint = context.add_checkpoint(**checkpoint_config)
+
+    assert context.checkpoint_store.save_count == 1
+
+    action_list = [
+        {
+            "name": "store_validation_result",
+            "action": {
+                "class_name": "StoreValidationResultAction",
+            },
+        },
+        {
+            "name": "store_evaluation_params",
+            "action": {
+                "class_name": "StoreEvaluationParametersAction",
+            },
+        },
+        {
+            "name": "update_data_docs",
+            "action": {
+                "class_name": "UpdateDataDocsAction",
+            },
+        },
+    ]
+    checkpoint.config.action_list = action_list
+    context.update_checkpoint(checkpoint)
+
+    assert context.checkpoint_store.save_count == 2
+
+
+@pytest.mark.unit
+def test_update_checkpoint_failure(in_memory_data_context: EphemeralDataContextSpy):
+    context = in_memory_data_context
+
+    name = "my_checkpoint"
+    checkpoint = Checkpoint(
+        name=name,
+        data_context=context,
+    )
+
+    with pytest.raises(gx_exceptions.CheckpointNotFoundError) as e:
+        context.update_checkpoint(checkpoint)
+
+    assert f"Could not find a Checkpoint named {name}" in str(e.value)
+
+
+@pytest.mark.unit
+def test_add_or_update_checkpoint_adds_successfully(
+    in_memory_data_context: EphemeralDataContextSpy,
+    checkpoint_config: dict,
+):
+    context = in_memory_data_context
+    name = "my_checkpoint"
+    checkpoint_config["name"] = name
+
+    checkpoint = context.add_or_update_checkpoint(**checkpoint_config)
+
+    actual_config = checkpoint.config
+
+    assert actual_config.name == name
+    assert (
+        actual_config.expectation_suite_name
+        == checkpoint_config["expectation_suite_name"]
+    )
+    assert actual_config.validations == checkpoint_config["validations"]
+    assert context.checkpoint_store.save_count == 1
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "update_kwargs,expected_config",
+    [
+        pytest.param(
+            {
+                "action_list": [],
+                "batch_request": {},
+                "class_name": "Checkpoint",
+                "config_version": 1.0,
+                "evaluation_parameters": {},
+                "expectation_suite_name": "oss_test_expectation_suite",
+                "module_name": "great_expectations.checkpoint",
+                "profilers": [],
+                "runtime_configuration": {},
+                "validations": [
+                    {"expectation_suite_name": "taxi.demo_pass"},
+                    {
+                        "batch_request": {
+                            "datasource_name": "oss_test_datasource",
+                            "data_connector_name": "oss_test_data_connector",
+                            "data_asset_name": "users",
+                        }
+                    },
+                ],
+            },
+            CheckpointConfig(
+                **{
+                    "action_list": [],
+                    "batch_request": {},
+                    "class_name": "Checkpoint",
+                    "config_version": 1.0,
+                    "evaluation_parameters": {},
+                    "expectation_suite_name": "oss_test_expectation_suite",
+                    "module_name": "great_expectations.checkpoint",
+                    "name": "my_checkpoint",
+                    "profilers": [],
+                    "runtime_configuration": {},
+                    "validations": [
+                        {"expectation_suite_name": "taxi.demo_pass"},
+                        {
+                            "batch_request": {
+                                "datasource_name": "oss_test_datasource",
+                                "data_connector_name": "oss_test_data_connector",
+                                "data_asset_name": "users",
+                            }
+                        },
+                    ],
+                }
+            ),
+            id="no updates",
+        ),
+        pytest.param(
+            {
+                "class_name": "Checkpoint",
+                "module_name": "great_expectations.checkpoint",
+                "config_version": 1.0,
+                "validations": [],
+                "evaluation_parameters": {
+                    "environment": "$GE_ENVIRONMENT",
+                    "tolerance": 1.0e-2,
+                    "aux_param_0": "$MY_PARAM",
+                    "aux_param_1": "1 + $MY_PARAM",
+                },
+            },
+            CheckpointConfig(
+                **{
+                    "action_list": [],
+                    "batch_request": {},
+                    "class_name": "Checkpoint",
+                    "config_version": 1.0,
+                    "evaluation_parameters": {
+                        "environment": "$GE_ENVIRONMENT",
+                        "tolerance": 1.0e-2,
+                        "aux_param_0": "$MY_PARAM",
+                        "aux_param_1": "1 + $MY_PARAM",
+                    },
+                    "expectation_suite_name": None,
+                    "module_name": "great_expectations.checkpoint",
+                    "name": "my_checkpoint",
+                    "profilers": [],
+                    "runtime_configuration": {},
+                    "validations": [],
+                }
+            ),
+            id="some updates",
+        ),
+    ],
+)
+def test_add_or_update_checkpoint_updates_successfully(
+    in_memory_data_context: EphemeralDataContextSpy,
+    checkpoint_config: dict,
+    update_kwargs: dict,
+    expected_config: CheckpointConfig,
+):
+    context = in_memory_data_context
+    name = "my_checkpoint"
+    checkpoint_config["name"] = name
+
+    checkpoint = context.add_checkpoint(**checkpoint_config)
+
+    assert context.checkpoint_store.save_count == 1
+
+    checkpoint = context.add_or_update_checkpoint(name=name, **update_kwargs)
+
+    assert checkpoint.config.to_dict() == expected_config.to_dict()
+    assert context.checkpoint_store.save_count == 2