Repository: great-expectations/great_expectations
Pull request #6991 (2023-01-30 16:35:30)

[MAINTENANCE] Standardize `RuleBasedProfiler` CRUD
Changes proposed in this pull request:
- Add `update_profiler` and `add_or_update_profiler`
- 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.


netlify[bot]: ### <span aria-hidden="true">👷</span> Deploy Preview for *niobium-lead-7998* processing.


|  Name | Link |
|---------------------------------|------------------------|
|<span aria-hidden="true">🔨</span> Latest commit | b913060f17bb911a46f4a03aaf7bbe95d09d9c00 |
|<span aria-hidden="true">🔍</span> Latest deploy log | https://app.netlify.com/sites/niobium-lead-7998/deploys/63d7f1d4ce939400091a4f06 |
codesee-maps[bot]: <details>
<summary><b>👇 <a href="https://app.codesee.io/r/reviews?pr=6991&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=6991&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/6991/209e71f2/b913060f17bb911a46f4a03aaf7bbe95d09d9c00.svg)](https://app.codesee.io/r/reviews?pr=6991&src=https%3A%2F%2Fgithub.com%2Fgreat-expectations%2Fgreat_expectations)

#### Legend
![CodeSee Map Legend](https://app.codesee.io/Legend_v2.png)

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 9b0fc59f3..591d77cd7 100644
--- a/great_expectations/data_context/data_context/abstract_data_context.py
+++ b/great_expectations/data_context/data_context/abstract_data_context.py
@@ -2486,6 +2486,61 @@ class AbstractDataContext(ConfigPeer, ABC):
             ge_cloud_id=ge_cloud_id,
         )
 
+    def update_profiler(self, profiler: RuleBasedProfiler) -> None:
+        """Update a Profiler that already exists.
+
+        Args:
+            profiler: The profiler to use to update.
+
+        Raises:
+            ProfilerNotFoundError: A profiler with the given name/id does not already exist.
+        """
+        RuleBasedProfiler.update_profiler(
+            profiler=profiler,
+            profiler_store=self.profiler_store,
+            data_context=self,
+        )
+
+    def add_or_update_profiler(
+        self,
+        name: str,
+        config_version: float | None = None,
+        rules: dict[str, dict] | None = None,
+        variables: dict | None = None,
+        id: str | None = None,
+    ) -> RuleBasedProfiler:
+        """Add a new Profiler or update an existing one on the context depending on whether it already exists or not.
+
+        Args:
+            name: The name of the RBP instance.
+            config_version: The version of the RBP (currently only 1.0 is supported).
+            rules: A set of dictionaries, each of which contains its own domain_builder, parameter_builders, and expectation_configuration_builders.
+            variables: Any variables to be substituted within the rules.
+            id: The id associated with the RBP instance (if applicable).
+
+        Returns:
+            A new Profiler or an updated one (depending on whether or not it existed before this method call).
+        """
+        existing_profiler = None
+        try:
+            existing_profiler = self.get_profiler(name=name, ge_cloud_id=id)
+        except gx_exceptions.ProfilerNotFoundError:
+            logger.info(
+                f"Could not find an existing profiler named '{name}'; creating a new one."
+            )
+
+        if existing_profiler:
+            self.update_profiler(existing_profiler)
+            return self.get_profiler(
+                name=existing_profiler.name, ge_cloud_id=existing_profiler.ge_cloud_id
+            )
+
+        config_version = config_version or 1.0
+        rules = rules or {}
+        return self.add_profiler(
+            name=name, config_version=config_version, rules=rules, variables=variables
+        )
+
     @usage_statistics_enabled_method(
         event_name=UsageStatsEvents.DATA_CONTEXT_RUN_RULE_BASED_PROFILER_WITH_DYNAMIC_ARGUMENTS,
     )
diff --git a/great_expectations/rule_based_profiler/rule_based_profiler.py b/great_expectations/rule_based_profiler/rule_based_profiler.py
index 4edb7fc7a..cfd9d50a9 100644
--- a/great_expectations/rule_based_profiler/rule_based_profiler.py
+++ b/great_expectations/rule_based_profiler/rule_based_profiler.py
@@ -1190,19 +1190,7 @@ class BaseRuleBasedProfiler(ConfigPeer):
         name: Optional[str] = None,
         ge_cloud_id: Optional[str] = None,
     ) -> RuleBasedProfiler:
-        assert bool(name) ^ bool(
-            ge_cloud_id
-        ), "Must provide either name or ge_cloud_id (but not both)"
-
-        key: Union[GXCloudIdentifier, ConfigurationIdentifier]
-        if ge_cloud_id:
-            key = GXCloudIdentifier(
-                resource_type=GXCloudRESTResource.PROFILER, cloud_id=ge_cloud_id
-            )
-        else:
-            key = ConfigurationIdentifier(
-                configuration_key=name,
-            )
+        key = RuleBasedProfiler._construct_profiler_key(name=name, id=ge_cloud_id)
         try:
             profiler_config: RuleBasedProfilerConfig = profiler_store.get(key=key)
         except gx_exceptions.InvalidKeyError as exc_ik:
@@ -1240,17 +1228,7 @@ class BaseRuleBasedProfiler(ConfigPeer):
         name: Optional[str] = None,
         ge_cloud_id: Optional[str] = None,
     ) -> None:
-        assert bool(name) ^ bool(
-            ge_cloud_id
-        ), "Must provide either name or ge_cloud_id (but not both)"
-
-        key: Union[GXCloudIdentifier, ConfigurationIdentifier]
-        if ge_cloud_id:
-            key = GXCloudIdentifier(
-                resource_type=GXCloudRESTResource.PROFILER, cloud_id=ge_cloud_id
-            )
-        else:
-            key = ConfigurationIdentifier(configuration_key=name)
+        key = RuleBasedProfiler._construct_profiler_key(name=name, id=ge_cloud_id)
 
         try:
             profiler_store.remove_key(key=key)
@@ -1264,6 +1242,38 @@ class BaseRuleBasedProfiler(ConfigPeer):
                 message=f'Non-existent Profiler configuration named "{id}".\n\nDetails: {exc_ik}'
             )
 
+    @staticmethod
+    def update_profiler(
+        profiler: RuleBasedProfiler,
+        profiler_store: ProfilerStore,
+        data_context: AbstractDataContext,
+    ) -> None:
+        name = profiler.name
+        id = profiler.ge_cloud_id
+
+        key = RuleBasedProfiler._construct_profiler_key(name=name, id=id)
+        if not profiler_store.has_key(key):
+            msg = "Non-existent Profiler configuration "
+            if name:
+                msg += f"named {name}"
+            elif id:
+                msg += f"with id {id}"
+            raise gx_exceptions.ProfilerNotFoundError(message=msg)
+
+        # Prioritize deletion by id if present (providing both will result in a downstream assertion error)
+        if id:
+            RuleBasedProfiler.delete_profiler(
+                ge_cloud_id=id, profiler_store=profiler_store
+            )
+        else:
+            RuleBasedProfiler.delete_profiler(name=name, profiler_store=profiler_store)
+
+        _ = RuleBasedProfiler.add_profiler(
+            config=profiler.config,
+            data_context=data_context,
+            profiler_store=profiler_store,
+        )
+
     @staticmethod
     def list_profilers(
         profiler_store: ProfilerStore,
@@ -1273,6 +1283,21 @@ class BaseRuleBasedProfiler(ConfigPeer):
             return profiler_store.list_keys()
         return [x.configuration_key for x in profiler_store.list_keys()]
 
+    @staticmethod
+    def _construct_profiler_key(
+        name: Optional[str] = None,
+        id: Optional[str] = None,
+    ) -> ConfigurationIdentifier | GXCloudIdentifier:
+        assert bool(name) ^ bool(id), "Must provide either name or id (but not both)"
+
+        if id:
+            return GXCloudIdentifier(
+                resource_type=GXCloudRESTResource.PROFILER, cloud_id=id
+            )
+        return ConfigurationIdentifier(
+            configuration_key=name,
+        )
+
     def self_check(self, pretty_print: bool = True) -> dict:
         """
         Necessary to enable integration with `AbstractDataContext.test_yaml_config`


Test patch:
diff --git a/tests/data_context/test_data_context_state_management.py b/tests/data_context/test_data_context_state_management.py
index 3cd6785d6..252a06467 100644
--- a/tests/data_context/test_data_context_state_management.py
+++ b/tests/data_context/test_data_context_state_management.py
@@ -10,7 +10,7 @@ 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.expectations_store import ExpectationsStore
+from great_expectations.data_context.store import ExpectationsStore, ProfilerStore
 from great_expectations.data_context.types.base import (
     DataContextConfig,
     DatasourceConfig,
@@ -18,6 +18,7 @@ from great_expectations.data_context.types.base import (
     ProgressBarsConfig,
 )
 from great_expectations.exceptions.exceptions import StoreConfigurationError
+from great_expectations.rule_based_profiler.rule_based_profiler import RuleBasedProfiler
 
 
 class ExpectationsStoreSpy(ExpectationsStore):
@@ -30,6 +31,19 @@ class ExpectationsStoreSpy(ExpectationsStore):
         return super().set(key=key, value=value, **kwargs)
 
 
+class ProfilerStoreSpy(ProfilerStore):
+
+    STORE_NAME = "profiler_store"
+
+    def __init__(self) -> None:
+        self.save_count = 0
+        super().__init__(ProfilerStoreSpy.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.
@@ -42,11 +56,16 @@ class EphemeralDataContextSpy(EphemeralDataContext):
         super().__init__(project_config)
         self.save_count = 0
         self._expectations_store = ExpectationsStoreSpy()
+        self._profiler_store = ProfilerStoreSpy()
 
     @property
     def expectations_store(self):
         return self._expectations_store
 
+    @property
+    def profiler_store(self):
+        return self._profiler_store
+
     def _save_project_config(self):
         """
         No-op our persistence mechanism but increment an internal counter to ensure it was used.
@@ -378,6 +397,91 @@ def test_add_or_update_expectation_suite_updates_successfully(
             kwargs={"column": "x", "value_set": [1, 2, 4]},
         ),
     ]
-    context.add_or_update_expectation_suite(expectation_suite_name=suite_name)
+    _ = context.add_or_update_expectation_suite(expectation_suite_name=suite_name)
 
     assert context.expectations_store.save_count == 2
+
+
+@pytest.mark.unit
+def test_update_profiler_success(
+    in_memory_data_context: EphemeralDataContextSpy,
+    profiler_rules: dict,
+):
+    context = in_memory_data_context
+
+    name = "my_rbp"
+    config_version = 1.0
+    rules = profiler_rules
+
+    profiler = context.add_profiler(
+        name=name, config_version=config_version, rules=rules
+    )
+
+    assert context.profiler_store.save_count == 1
+
+    profiler.rules = {}
+    context.update_profiler(profiler)
+
+    assert context.profiler_store.save_count == 2
+
+
+@pytest.mark.unit
+def test_update_profiler_failure(in_memory_data_context: EphemeralDataContextSpy):
+    context = in_memory_data_context
+
+    name = "my_rbp"
+    profiler = RuleBasedProfiler(
+        name=name,
+        config_version=1.0,
+        data_context=context,
+    )
+
+    with pytest.raises(gx_exceptions.ProfilerNotFoundError) as e:
+        context.update_profiler(profiler)
+
+    assert f"Non-existent Profiler configuration named {name}" in str(e.value)
+
+
+@pytest.mark.unit
+def test_add_or_update_profiler_adds_successfully(
+    in_memory_data_context: EphemeralDataContextSpy, profiler_rules: dict
+):
+    context = in_memory_data_context
+
+    name = "my_rbp"
+    config_version = 1.0
+    rules = profiler_rules
+
+    profiler = context.add_or_update_profiler(
+        name=name, config_version=config_version, rules=rules
+    )
+
+    config = profiler.config
+
+    assert config.name == name
+    assert config.config_version == config_version
+    assert len(config.rules) == len(rules) and config.rules.keys() == rules.keys()
+    assert context.profiler_store.save_count == 1
+
+
+@pytest.mark.unit
+def test_add_or_update_profiler_updates_successfully(
+    in_memory_data_context: EphemeralDataContextSpy,
+    profiler_rules: dict,
+):
+    context = in_memory_data_context
+
+    name = "my_rbp"
+    config_version = 1.0
+    rules = profiler_rules
+
+    profiler = context.add_profiler(
+        name=name, config_version=config_version, rules=rules
+    )
+
+    assert context.profiler_store.save_count == 1
+
+    profiler.rules = {}
+    _ = context.add_or_update_profiler(name=name)
+
+    assert context.profiler_store.save_count == 2
diff --git a/tests/rule_based_profiler/test_rule_based_profiler.py b/tests/rule_based_profiler/test_rule_based_profiler.py
index a6744d57f..e6ce720f7 100644
--- a/tests/rule_based_profiler/test_rule_based_profiler.py
+++ b/tests/rule_based_profiler/test_rule_based_profiler.py
@@ -1098,7 +1098,7 @@ def test_get_profiler_with_too_many_args_raises_error(
             ge_cloud_id="my_ge_cloud_id",
         )
 
-    assert "either name or ge_cloud_id" in str(e.value)
+    assert "either name or id" in str(e.value)
 
 
 @mock.patch("great_expectations.data_context.data_context.AbstractDataContext")
@@ -1320,7 +1320,7 @@ def test_delete_profiler_with_too_many_args_raises_error(
             ge_cloud_id="my_ge_cloud_id",
         )
 
-    assert "either name or ge_cloud_id" in str(e.value)
+    assert "either name or id" in str(e.value)
 
 
 @pytest.mark.unit