Repository: great-expectations/great_expectations
Pull request #7060 (2023-02-06 19:32:25)

[FEATURE] Add `profiler` arg to `DataContext` `Profiler` CRUD
Changes proposed in this pull request:
- Allow users to pass in an existing profiler through `add_profiler` and `add_or_update_profiler`


### 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 made corresponding changes to the documentation
- [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 | 78c5a3c183006ad5a59fd4848e182400c39ce056 |
|<span aria-hidden="true">🔍</span> Latest deploy log | https://app.netlify.com/sites/niobium-lead-7998/deploys/63e155cb143cc60008044b9b |
codesee-maps[bot]: <details>
<summary><b>👇 <a href="https://app.codesee.io/r/reviews?pr=7060&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=7060&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/7060/bd3abe0d/78c5a3c183006ad5a59fd4848e182400c39ce056.svg)](https://app.codesee.io/r/reviews?pr=7060&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 f8c808e6a..9684af8d2 100644
--- a/great_expectations/data_context/data_context/abstract_data_context.py
+++ b/great_expectations/data_context/data_context/abstract_data_context.py
@@ -2943,23 +2943,60 @@ class AbstractDataContext(ConfigPeer, ABC):
                 f"expectation_suite {expectation_suite_name} not found"
             )
 
-    @public_api
+    @overload
     def add_profiler(
         self,
         name: str,
         config_version: float,
-        rules: Dict[str, dict],
-        variables: Optional[dict] = None,
+        rules: dict[str, dict],
+        variables: dict | None = ...,
+        profiler: None = ...,
+    ) -> RuleBasedProfiler:
+        """
+        Individual constructors args (`name`, `config_version`, and `rules`) are provided.
+        `profiler` should not be provided.
+        """
+        ...
+
+    @overload
+    def add_profiler(
+        self,
+        name: None = ...,
+        config_version: None = ...,
+        rules: None = ...,
+        variables: None = ...,
+        profiler: RuleBasedProfiler = ...,
+    ) -> RuleBasedProfiler:
+        """
+        `profiler` is provided.
+        Individual constructors args (`name`, `config_version`, and `rules`) should not be provided.
+        """
+        ...
+
+    @public_api
+    @new_argument(
+        argument_name="profiler",
+        version="0.15.48",
+        message="Pass in an existing profiler instead of individual constructor args",
+    )
+    def add_profiler(
+        self,
+        name: str | None = None,
+        config_version: float | None = None,
+        rules: dict[str, dict] | None = None,
+        variables: dict | None = None,
+        profiler: RuleBasedProfiler | None = None,
     ) -> RuleBasedProfiler:
         """
         Constructs a Profiler, persists it utilizing the context's underlying ProfilerStore,
         and returns it to the user for subsequent usage.
 
         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
-            variables: Any variables to be substituted within the rules
+            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.
+            profiler: An existing RuleBasedProfiler to persist.
 
         Returns:
             The persisted Profiler constructed by the input arguments.
@@ -2970,31 +3007,44 @@ class AbstractDataContext(ConfigPeer, ABC):
             config_version=config_version,
             rules=rules,
             variables=variables,
+            profiler=profiler,
         )
 
     def _add_profiler(
         self,
-        name: str,
+        name: str | None,
         id: str | None,
-        config_version: float,
-        rules: dict[str, dict],
+        config_version: float | None,
+        rules: dict[str, dict] | None,
         variables: dict | None,
+        profiler: RuleBasedProfiler | None,
     ) -> RuleBasedProfiler:
-        config_data = {
-            "name": name,
-            "id": id,
-            "config_version": config_version,
-            "rules": rules,
-            "variables": variables,
-        }
+        if not (
+            (profiler is None)
+            ^ all(arg is None for arg in (name, config_version, rules))
+        ):
+            raise ValueError(
+                "Must either pass in an existing profiler or individual constructor arguments (but not both)"
+            )
 
-        # Roundtrip through schema validation to remove any illegal fields add/or restore any missing fields.
-        validated_config: dict = ruleBasedProfilerConfigSchema.load(config_data)
-        profiler_config: dict = ruleBasedProfilerConfigSchema.dump(validated_config)
-        profiler_config.pop("class_name")
-        profiler_config.pop("module_name")
+        if profiler:
+            config = profiler.config
+        else:
+            config_data = {
+                "name": name,
+                "id": id,
+                "config_version": config_version,
+                "rules": rules,
+                "variables": variables,
+            }
 
-        config = RuleBasedProfilerConfig(**profiler_config)
+            # Roundtrip through schema validation to remove any illegal fields add/or restore any missing fields.
+            validated_config: dict = ruleBasedProfilerConfigSchema.load(config_data)
+            profiler_config: dict = ruleBasedProfilerConfigSchema.dump(validated_config)
+            profiler_config.pop("class_name")
+            profiler_config.pop("module_name")
+
+            config = RuleBasedProfilerConfig(**profiler_config)
 
         profiler = RuleBasedProfiler.add_profiler(
             config=config,
@@ -3088,15 +3138,46 @@ class AbstractDataContext(ConfigPeer, ABC):
             data_context=self,
         )
 
+    @overload
+    def add_or_update_profiler(
+        self,
+        name: str,
+        config_version: float,
+        rules: dict[str, dict],
+        variables: dict | None = ...,
+        profiler: None = ...,
+    ) -> RuleBasedProfiler:
+        """
+        Individual constructors args (`name`, `config_version`, and `rules`) are provided.
+        `profiler` should not be provided.
+        """
+        ...
+
+    @overload
+    def add_or_update_profiler(
+        self,
+        name: None = ...,
+        config_version: None = ...,
+        rules: None = ...,
+        variables: None = ...,
+        profiler: RuleBasedProfiler = ...,
+    ) -> RuleBasedProfiler:
+        """
+        `profiler` is provided.
+        Individual constructors args (`name`, `config_version`, and `rules`) should not be provided.
+        """
+        ...
+
     @public_api
     @new_method_or_class(version="0.15.48")
     def add_or_update_profiler(
         self,
-        name: str,
+        name: str | None = None,
         id: str | None = None,
         config_version: float | None = None,
         rules: dict[str, dict] | None = None,
         variables: dict | None = None,
+        profiler: RuleBasedProfiler | None = None,
     ) -> RuleBasedProfiler:
         """Add a new Profiler or update an existing one on the context depending on whether it already exists or not.
 
@@ -3106,18 +3187,18 @@ class AbstractDataContext(ConfigPeer, ABC):
             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).
+            profiler: An existing RuleBasedProfiler to persist.
 
         Returns:
             A new Profiler or an updated one (depending on whether or not it existed before this method call).
         """
-        config_version = config_version or 1.0
-        rules = rules or {}
         return self._add_profiler(
             name=name,
             id=id,
             config_version=config_version,
             rules=rules,
             variables=variables,
+            profiler=profiler,
         )
 
     @usage_statistics_enabled_method(


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 7d0d63013..e95029412 100644
--- a/tests/data_context/test_data_context_state_management.py
+++ b/tests/data_context/test_data_context_state_management.py
@@ -566,6 +566,58 @@ def test_add_or_update_expectation_suite_conflicting_args_failure(
     assert context.expectations_store.save_count == 0
 
 
+@pytest.mark.unit
+def test_add_profiler_with_existing_profiler(
+    in_memory_data_context: EphemeralDataContextSpy,
+    profiler_rules: dict,
+):
+    context = in_memory_data_context
+
+    name = "my_rbp"
+    profiler = RuleBasedProfiler(
+        name=name,
+        config_version=1.0,
+        rules=profiler_rules,
+        data_context=context,
+    )
+
+    persisted_profiler = context.add_profiler(profiler=profiler)
+
+    assert profiler.name == persisted_profiler.name
+    assert profiler.config_version == persisted_profiler.config_version
+    assert len(profiler.rules) == len(persisted_profiler.rules)
+    assert context.profiler_store.save_count == 1
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "profiler,profiler_name",
+    [
+        pytest.param(
+            mock.MagicMock(),
+            "my_rbp",
+            id="both profiler and profiler_name",
+        ),
+        pytest.param(None, None, id="neither profiler nor profiler_name"),
+    ],
+)
+def test_add_profiler_conflicting_args_failure(
+    in_memory_data_context: EphemeralDataContextSpy,
+    profiler: mock.MagicMock | None,
+    profiler_name: str | None,
+):
+    context = in_memory_data_context
+
+    with pytest.raises(ValueError) as e:
+        context.add_profiler(profiler=profiler, name=profiler_name)
+
+    assert (
+        "an existing profiler or individual constructor arguments (but not both)"
+        in str(e.value)
+    )
+    assert context.profiler_store.save_count == 0
+
+
 @pytest.mark.unit
 def test_update_profiler_success(
     in_memory_data_context: EphemeralDataContextSpy,
@@ -682,13 +734,44 @@ def test_add_or_update_profiler_updates_successfully(
 
     assert context.profiler_store.save_count == 1
 
-    profiler = context.add_or_update_profiler(name=name, rules=new_rules)
+    profiler = context.add_or_update_profiler(
+        name=name, rules=new_rules, config_version=config_version
+    )
 
     # 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
+@pytest.mark.parametrize(
+    "profiler,profiler_name",
+    [
+        pytest.param(
+            mock.MagicMock(),  # Only care about the presence of the value (no need to construct a full RBP obj)
+            "my_rbp",
+            id="both profiler and profiler_name",
+        ),
+        pytest.param(None, None, id="neither profiler nor profiler_name"),
+    ],
+)
+def test_add_or_update_profiler_conflicting_args_failure(
+    in_memory_data_context: EphemeralDataContextSpy,
+    profiler: mock.MagicMock | None,
+    profiler_name: str | None,
+):
+    context = in_memory_data_context
+
+    with pytest.raises(ValueError) as e:
+        context.add_or_update_profiler(profiler=profiler, name=profiler_name)
+
+    assert (
+        "an existing profiler or individual constructor arguments (but not both)"
+        in str(e.value)
+    )
+    assert context.profiler_store.save_count == 0
+
+
 @pytest.mark.unit
 def test_add_checkpoint_with_existing_checkpoint(
     in_memory_data_context: EphemeralDataContextSpy,