[FEATURE] CheckpointFactory.add_or_update

great_expectations/core/factory/checkpoint_factory.py

@@ -12,6 +12,7 @@
 from great_expectations.checkpoint.checkpoint import Checkpoint
 from great_expectations.compatibility.typing_extensions import override
 from great_expectations.core.factory.factory import Factory
+from great_expectations.data_context import project_manager
 from great_expectations.exceptions import DataContextError
 
 if TYPE_CHECKING:
@@ -129,3 +130,33 @@ def _get(self, key: GXCloudIdentifier | StringKey) -> Checkpoint:
             raise ValueError(f"Object with key {key} was found, but it is not a Checkpoint.")  # noqa: TRY003, TRY004 # FIXME CoP
 
         return checkpoint
+
+    @public_api
+    def add_or_update(self, checkpoint: Checkpoint) -> Checkpoint:
+        """Add or update a Checkpoint by name.
+
+        If an Checkpoint with the same name exists, overwrite it, otherwise
+        create a new Checkpoint. Calls add_or_update
+
+        Args:
+            checkpoint: Checkpoint to add or update
+        """
+        try:
+            existing_checkpoint = self.get(name=checkpoint.name)
+        except DataContextError:
+            # checkpoint doesn't exist yet, so add it
+            return self.add(checkpoint=checkpoint)
+
+        # update checkpoint
+        checkpoint.id = existing_checkpoint.id
+
+        val_def_ids_by_name = {
+            val_def.name: val_def.id for val_def in existing_checkpoint.validation_definitions
+        }
+        val_def_factory = project_manager.get_validation_definitions_factory()
+        for val_def in checkpoint.validation_definitions:
+            if val_def.name in val_def_ids_by_name:
+                val_def.id = val_def_ids_by_name[val_def.name]
+            val_def_factory.add_or_update(validation_definition=val_def)
+        checkpoint.save()
+        return checkpoint

great_expectations/core/factory/factory.py

@@ -24,3 +24,7 @@ def get(self, name: str) -> T:
     @abstractmethod
     def all(self) -> Iterable[T]:
         pass
+
+    @abstractmethod
+    def add_or_update(self, obj: T) -> T:
+        pass

great_expectations/data_context/data_context/context_factory.py

@@ -15,6 +15,7 @@
 
 import great_expectations.exceptions as gx_exceptions
 from great_expectations._docs_decorators import public_api
+from great_expectations.core.factory import SuiteFactory, ValidationDefinitionFactory
 from great_expectations.exceptions import (
     GXCloudConfigurationError,
 )
@@ -97,6 +98,9 @@
     def get_expectations_store(self) -> ExpectationsStore:
         return self._project.expectations_store
 
+    def get_suites_factory(self) -> SuiteFactory:
+        return self.__project.suites
+
     def get_checkpoints_store(self) -> CheckpointStore:
         return self._project.checkpoint_store
 
@@ -106,6 +110,9 @@
     def get_validation_definition_store(self) -> ValidationDefinitionStore:
         return self._project.validation_definition_store
 
+    def get_validation_definitions_factory(self) -> ValidationDefinitionFactory:
+        return self.__project.validation_definitions
+
     def get_datasources(self) -> DatasourceDict:
         return self._project.data_sources.all()
 

tests/core/factory/test_checkpoint_factory.py

@@ -1,5 +1,6 @@
 import re
 from unittest import mock
+from unittest.mock import ANY
 from unittest.mock import ANY as ANY_TEST_ARG
 
 import pytest
@@ -415,3 +416,152 @@ def _test_checkpoint_factory_delete_emits_event(self, context):
         mock_submit.assert_called_once_with(
             event=CheckpointDeletedEvent(checkpoint_id=checkpoint.id)
         )
+
+
+class TestCheckpointFactoryAddOrUpdate:
+    @pytest.mark.filesystem
+    def test_add_empty_new_checkpoint__filesystem(self, empty_data_context):
+        return self._test_add_empty_new_checkpoint(empty_data_context)
+
+    @pytest.mark.cloud
+    def test_add_empty_new_checkpoint__cloud(self, empty_cloud_context_fluent):
+        return self._test_add_empty_new_checkpoint(empty_cloud_context_fluent)
+
+    @pytest.mark.unit
+    def test_add_empty_new_checkpoint__ephemeral(self, ephemeral_context_with_defaults):
+        return self._test_add_empty_new_checkpoint(ephemeral_context_with_defaults)
+
+    def _test_add_empty_new_checkpoint(self, context: AbstractDataContext):
+        # arrange
+        checkpoint_name = "checkpoint A"
+        checkpoint = Checkpoint(name=checkpoint_name, validation_definitions=[])
+
+        # act
+        created_checkpoint = context.checkpoints.add_or_update(checkpoint=checkpoint)
+
+        # assert
+        assert created_checkpoint.id
+        context.checkpoints.get(checkpoint_name)
+
+    @pytest.mark.filesystem
+    def test_add_new_checkpoint_with_validations__filesystem(self, empty_data_context):
+        return self._test_add_new_checkpoint_with_validations(empty_data_context)
+
+    @pytest.mark.cloud
+    def test_add_new_checkpoint_with_validations__cloud(self, empty_cloud_context_fluent):
+        return self._test_add_new_checkpoint_with_validations(empty_cloud_context_fluent)
+
+    @pytest.mark.unit
+    def test_add_new_checkpoint_with_validations__ephemeral(self, ephemeral_context_with_defaults):
+        return self._test_add_new_checkpoint_with_validations(ephemeral_context_with_defaults)
+
+    def _test_add_new_checkpoint_with_validations(self, context: AbstractDataContext):
+        # arrange
+        checkpoint_name = "checkpoint A"
+        batch_def = (
+            context.data_sources.add_pandas("data source A")
+            .add_dataframe_asset("asset A")
+            .add_batch_definition_whole_dataframe("batch def A")
+        )
+
+        validation_definitions = [
+            ValidationDefinition(
+                name="val def A",
+                data=batch_def,
+                suite=ExpectationSuite(name="suite A"),
+            ),
+            ValidationDefinition(
+                name="val def B",
+                data=batch_def,
+                suite=ExpectationSuite(name="suite B"),
+            ),
+        ]
+        checkpoint = Checkpoint(name=checkpoint_name, validation_definitions=validation_definitions)
+
+        # act
+        created_checkpoint = context.checkpoints.add_or_update(checkpoint=checkpoint)
+
+        # assert
+        assert created_checkpoint.id
+        for val_def, created_val_def in zip(
+            validation_definitions, created_checkpoint.validation_definitions
+        ):
+            assert created_val_def.id
+            val_def.id = ANY
+            assert val_def == created_val_def
+
+    @pytest.mark.filesystem
+    def test_update_existing_checkpoint_adds_validations__filesystem(self, empty_data_context):
+        return self._test_update_existing_checkpoint_adds_validations(empty_data_context)
+
+    @pytest.mark.cloud
+    def test_update_existing_checkpoint_adds_validations__cloud(self, empty_cloud_context_fluent):
+        return self._test_update_existing_checkpoint_adds_validations(empty_cloud_context_fluent)
+
+    @pytest.mark.unit
+    def test_update_existing_checkpoint_adds_validations__ephemeral(
+        self, ephemeral_context_with_defaults
+    ):
+        return self._test_update_existing_checkpoint_adds_validations(
+            ephemeral_context_with_defaults
+        )
+
+    def _test_update_existing_checkpoint_adds_validations(self, context: AbstractDataContext): ...
+
+    @pytest.mark.filesystem
+    def test_update_existing_checkpoint_updates_validations__filesystem(self, empty_data_context):
+        return self._test_update_existing_checkpoint_updates_validations(empty_data_context)
+
+    @pytest.mark.cloud
+    def test_update_existing_checkpoint_updates_validations__cloud(
+        self, empty_cloud_context_fluent
+    ):
+        return self._test_update_existing_checkpoint_updates_validations(empty_cloud_context_fluent)
+
+    @pytest.mark.unit
+    def test_update_existing_checkpoint_updates_validations__ephemeral(
+        self, ephemeral_context_with_defaults
+    ):
+        return self._test_update_existing_checkpoint_updates_validations(
+            ephemeral_context_with_defaults
+        )
+
+    def _test_update_existing_checkpoint_updates_validations(
+        self, context: AbstractDataContext
+    ): ...
+
+    @pytest.mark.filesystem
+    def test_update_existing_checkpoint_deletes_validations__filesystem(self, empty_data_context):
+        return self._test_update_existing_checkpoint_deletes_validations(empty_data_context)
+
+    @pytest.mark.cloud
+    def test_update_existing_checkpoint_deletes_validations__cloud(
+        self, empty_cloud_context_fluent
+    ):
+        return self._test_update_existing_checkpoint_deletes_validations(empty_cloud_context_fluent)
+
+    @pytest.mark.unit
+    def test_update_existing_checkpoint_deletes_validations__ephemeral(
+        self, ephemeral_context_with_defaults
+    ):
+        return self._test_update_existing_checkpoint_deletes_validations(
+            ephemeral_context_with_defaults
+        )
+
+    def _test_update_existing_checkpoint_deletes_validations(
+        self, context: AbstractDataContext
+    ): ...
+
+    @pytest.mark.filesystem
+    def test_add_or_update_is_idempotent__filesystem(self, empty_data_context):
+        return self._test_add_or_update_is_idempotent(empty_data_context)
+
+    @pytest.mark.cloud
+    def test_add_or_update_is_idempotent__cloud(self, empty_cloud_context_fluent):
+        return self._test_add_or_update_is_idempotent(empty_cloud_context_fluent)
+
+    @pytest.mark.unit
+    def test_add_or_update_is_idempotent__ephemeral(self, ephemeral_context_with_defaults):
+        return self._test_add_or_update_is_idempotent(ephemeral_context_with_defaults)
+
+    def _test_add_or_update_is_idempotent(self, context: AbstractDataContext): ...