Repository: PrefectHQ/prefect

Issue:
Improve or create new GE task library task to run a checkpoint against in-memory config
## Current behavior
Today there is a [great expectations task in the task library to run GE checkpoints](https://docs.prefect.io/api/latest/tasks/great_expectations.html#prefect-tasks-great_expectations-checkpoint-rungreatexpectationscheckpoint), but it is dependent on being run in the correct environment to find the local data configuration files specific to the expectation suite being run, so it is not super great in a distributed environment. In addition, it doesn't meaningfully utilize the dataflow semantics of Prefect since the GE checkpoint infrastructure re-reads the data into memory based on config of where it is stored (disk/sql/spark).


## Proposed behavior
Preferably, you should be able to pass all the relevant configuration to run a GE checkpoint as Python objects in-memory so we don't have to depend on the local filesystem matching up on a given worker, including
1. any config GE needs to know about the expectation suite store
2. any config GE needs to know about *which* actions it should take (including whether to store to a validation store or rebuild data docs)
3. any config GE needs to know to *actually do* the actions it should take (such as where the validation store IS and how to connect to it)
4. the actual data object to validate in memory


After talking some with the GE folks (cc @jcampbell, @abegong) there is a way to configure your [expectations suites](https://docs.greatexpectations.io/en/latest/how_to_guides/configuring_metadata_stores/how_to_configure_an_expectation_store_in_amazon_s3.html) and [validation store](https://docs.greatexpectations.io/en/latest/how_to_guides/configuring_metadata_stores/how_to_configure_a_validation_result_store_in_s3.html) in S3. In regular usage, you still need a `great_expectations.yml` file that points to those locations and to configure the data contexts that you are actually running those expectations against. BUT also from talking with them it sounds that by instantiating the [BaseDataContext class](https://github.com/great-expectations/great_expectations/blob/4f2e634502d81bbeaf5e719b8e6611f26a12a4a4/great_expectations/data_context/data_context.py#L84) directly, we should be able to pass a config object in memory that can be composed of all the relevant stuff we need from `great_expectations.yml` (we will just have to reverse engineer it a bit I think because I can't find any good API docs on what that config object should look like). According to them this class was originally made to make testing easier to do in memory, so there should be some inspiration in the test files (in particular [here is a datacontextconfig](https://github.com/great-expectations/great_expectations/blob/4f2e634502d81bbeaf5e719b8e6611f26a12a4a4/tests/data_context/test_data_context.py#L591) used for testing that should be good inspo for what we need to cover from `great_expectations.yml` in memory).

Finally, the last piece to make this most Prefect-y is probably to pass the actual piece of data as an in-memory object instead of as a lot of configuration around files/sql/spark dfs. I don't have a great idea of how to do this yet so this may need to be punted to another issue. I had originally tried to do this by instantiating batches and passing them directly into [`DataContext.run_validation_operator.assets_to_validate`](https://github.com/PrefectHQ/prefect/blob/master/src/prefect/tasks/great_expectations/checkpoints.py#L110) but those actually must be [`DataAsset`s or arguments to get_batch](https://github.com/great-expectations/great_expectations/commit/9a31ff2b0596c57d02fa2c36919e86a12b493e49) which will be re-read by GE infra and which also need a ton of config related not just to the data but the expectation suite.


## Example
<!-- Please give an example of how the enhancement would be useful -->
- Users who want to do GE validation on distributed (and especially ephemeral) infrastructure that would be annoying to keep in sync with a local file system that has the GE directories on it


Discussion:

Resolving patch:
diff --git a/src/prefect/tasks/great_expectations/checkpoints.py b/src/prefect/tasks/great_expectations/checkpoints.py
--- a/src/prefect/tasks/great_expectations/checkpoints.py
+++ b/src/prefect/tasks/great_expectations/checkpoints.py
@@ -10,10 +10,11 @@
 """
 import prefect
 from prefect import Task
+from prefect.artifacts import create_markdown
+
 from prefect.engine import signals
 from prefect.utilities.tasks import defaults_from_attrs
 
-
 import great_expectations as ge
 
 from typing import Optional
@@ -21,97 +22,237 @@
 
 class RunGreatExpectationsCheckpoint(Task):
     """
-    Task for running a Great Expectations checkpoint. For this task to run properly, it must be
-    run above your great_expectations directory or configured with the `context_root_dir` for
-    your great_expectations directory on the local file system of the worker process.
+    Task for running a Great Expectations checkpoint.
+
+    Example using the GE getting started tutorial:
+    https://github.com/superconductive/ge_tutorials/tree/main/ge_getting_started_tutorial
+
+    This is a parameterized flow for building the `batch_kwargs` to be passed into the task.
+
+    ```python
+    from prefect import task, Flow, Parameter
+    from prefect.tasks.great_expectations import RunGreatExpectationsCheckpoint
+
+
+    # Define checkpoint task
+    checkpoint_task = RunGreatExpectationsCheckpoint()
+
+
+    # Task for retrieving batch kwargs including csv dataset
+    @task
+    def get_batch_kwargs(datasource_name, dataset):
+        dataset = ge.read_csv(dataset)
+        return {"dataset": dataset, "datasource": datasource_name}
+
+
+    with Flow("ge_test") as flow:
+        datasource_name = Parameter("datasource_name")
+        dataset = Parameter("dataset")
+        batch_kwargs = get_batch_kwargs(datasource_name, dataset)
+
+        expectation_suite_name = Parameter("expectation_suite_name")
+        checkpoint_name = Parameter("checkpoint_name")
+        checkpoint_task(
+            checkpoint_name=checkpoint_name,
+            batch_kwargs=batch_kwargs,
+            expectation_suite_name=expectation_suite_name,
+        )
+
+    flow.run(
+        parameters={
+            "datasource_name": "data__dir",
+            "dataset": "data/yellow_tripdata_sample_2019-01.csv",
+            "expectation_suite_name": "yellow_tripdata_sample_2019-01.warning",
+            "checkpoint_name": "action_list_operator",
+        },
+    )
+    ```
+
 
     Args:
-        - checkpoint_name (str): the name of the checkpoint; should match the filename of the
-            checkpoint without .py
-        - context_root_dir (str): the absolute or relative path to the directory holding your
-            `great_expectations.yml`
-        - runtime_environment (dict): a dictionary of great expectation config key-value pairs
-            to overwrite your config in `great_expectations.yml`
-        - run_name (str): the name of this Great Expectation validation run; defaults to the
-            task slug
+        - checkpoint_name (str, optional): the name of the checkpoint; should match the filename of
+            the checkpoint without .py
+        - context (DataContext, optional): an in-memory GE DataContext object. e.g.
+            `ge.data_context.DataContext()` If not provided then `context_root_dir` will be used to
+            look for one.
+        - assets_to_validate (list, optional): A list of assets to validate when running the
+            validation operator. If not provided then `batch_kwargs` and `expectation_suite_name`
+            will be used if context is provided. Also, if not provided and
+            `get_checkpoint_from_context` is True then the assets will be loaded from that context.
+        - batch_kwargs (dict, optional): a dictionary of batch kwargs to be used when validating
+            assets.
+        - expectation_suite_name (str, optional): the name of an expectation suite to be used when
+            validating assets.
+        - get_checkpoint_from_context (bool, optional): get the checkpoint from context. Defaults to
+            `False`
+        - context_root_dir (str, optional): the absolute or relative path to the directory holding
+            your `great_expectations.yml`
+        - runtime_environment (dict, optional): a dictionary of great expectation config key-value
+            pairs to overwrite your config in `great_expectations.yml`
+        - run_name (str, optional): the name of this  Great Expectation validation run; defaults to
+            the task slug
+        - run_info_at_end (bool, optional): add run info to the end of the artifact generated by this
+            task. Defaults to `True`.
+        - disable_markdown_artifact (bool, optional): toggle the posting of a markdown artifact from
+            this tasks. Defaults to `False`.
         - **kwargs (dict, optional): additional keyword arguments to pass to the Task constructor
     """
 
     def __init__(
         self,
         checkpoint_name: str = None,
+        context: "ge.DataContext" = None,
+        assets_to_validate: list = None,
+        batch_kwargs: dict = None,
+        expectation_suite_name: str = None,
+        get_checkpoint_from_context: bool = False,
         context_root_dir: str = None,
         runtime_environment: Optional[dict] = None,
         run_name: str = None,
+        run_info_at_end: bool = True,
+        disable_markdown_artifact: bool = False,
         **kwargs
     ):
         self.checkpoint_name = checkpoint_name
+        self.context = context
+        self.assets_to_validate = assets_to_validate
+        self.batch_kwargs = batch_kwargs
+        self.expectation_suite_name = expectation_suite_name
+        self.get_checkpoint_from_context = get_checkpoint_from_context
         self.context_root_dir = context_root_dir
         self.runtime_environment = runtime_environment or dict()
         self.run_name = run_name
+        self.run_info_at_end = run_info_at_end
+        self.disable_markdown_artifact = disable_markdown_artifact
 
         super().__init__(**kwargs)
 
     @defaults_from_attrs(
-        "checkpoint_name", "context_root_dir", "runtime_environment", "run_name"
+        "checkpoint_name",
+        "context",
+        "assets_to_validate",
+        "batch_kwargs",
+        "expectation_suite_name",
+        "get_checkpoint_from_context",
+        "context_root_dir",
+        "runtime_environment",
+        "run_name",
+        "run_info_at_end",
+        "disable_markdown_artifact",
     )
     def run(
         self,
         checkpoint_name: str = None,
+        context: "ge.DataContext" = None,
+        assets_to_validate: list = None,
+        batch_kwargs: dict = None,
+        expectation_suite_name: str = None,
+        get_checkpoint_from_context: bool = False,
         context_root_dir: str = None,
         runtime_environment: Optional[dict] = None,
         run_name: str = None,
-        **kwargs
+        run_info_at_end: bool = True,
+        disable_markdown_artifact: bool = False,
     ):
         """
         Task run method.
 
         Args:
-            - checkpoint_name (str): the name of the checkpoint; should match the filename of
+            - checkpoint_name (str, optional): the name of the checkpoint; should match the filename of
                 the checkpoint without .py
-            - context_root_dir (str): the absolute or relative path to the directory holding
+            - context (DataContext, optional): an in-memory GE DataContext object. e.g.
+                `ge.data_context.DataContext()` If not provided then `context_root_dir` will be used to
+                look for one.
+            - assets_to_validate (list, optional): A list of assets to validate when running the
+                validation operator. If not provided then `batch_kwargs` and `expectation_suite_name`
+                will be used if context is provided. Also, if not provided and
+                `get_checkpoint_from_context` is True then the assets will be loaded from that context.
+            - batch_kwargs (dict, optional): a dictionary of batch kwargs to be used when validating
+                assets.
+            - expectation_suite_name (str, optional): the name of an expectation suite to be used when
+                validating assets.
+            - get_checkpoint_from_context (bool, optional): get the checkpoint from context. Defaults to
+                `False`
+            - context_root_dir (str, optional): the absolute or relative path to the directory holding
                 your `great_expectations.yml`
-            - runtime_environment (dict): a dictionary of great expectation config key-value
+            - runtime_environment (dict, optional): a dictionary of great expectation config key-value
                 pairs to overwrite your config in `great_expectations.yml`
-            - run_name (str): the name of this  Great Expectation validation run; defaults to
+            - run_name (str, optional): the name of this  Great Expectation validation run; defaults to
                 the task slug
-            - **kwargs (dict, optional): additional keyword arguments to pass to the Task
-                constructor
+            - run_info_at_end (bool, optional): add run info to the end of the artifact generated by this
+                task. Defaults to `True`.
+            - disable_markdown_artifact (bool, optional): toggle the posting of a markdown artifact from
+                this tasks. Defaults to `False`.
 
         Raises:
             - 'signals.VALIDATIONFAIL' if the validation was not a success
+
         Returns:
             - result
                 ('great_expectations.validation_operators.types.validation_operator_result.ValidationOperatorResult'):
                 The Great Expectations metadata returned from the validation
 
         """
-
         if checkpoint_name is None:
             raise ValueError("You must provide the checkpoint name.")
 
         runtime_environment = runtime_environment or dict()
 
-        context = ge.DataContext(
-            context_root_dir=context_root_dir, runtime_environment=runtime_environment
-        )
-        checkpoint = context.get_checkpoint(checkpoint_name)
+        # Load context if not provided directly
+        if not context:
+            context = ge.DataContext(
+                context_root_dir=context_root_dir,
+                runtime_environment=runtime_environment,
+            )
+
+        # if assets are not provided directly through `assets_to_validate` then they need be loaded
+        #   if the checkpoint is being loaded from the context then load suite and batch from there
+        #   otherwise get batch from `batch_kwargs` and `expectation_suite_name`
+        if not assets_to_validate:
+            assets_to_validate = []
+            if get_checkpoint_from_context:
+                ge_checkpoint = context.get_checkpoint(checkpoint_name)
 
-        batches_to_validate = []
-        for batch in checkpoint["batches"]:
-            batch_kwargs = batch["batch_kwargs"]
-            for suite_name in batch["expectation_suite_names"]:
-                suite = context.get_expectation_suite(suite_name)
-                batch = context.get_batch(batch_kwargs, suite)
-                batches_to_validate.append(batch)
+                for batch in ge_checkpoint["batches"]:
+                    batch_kwargs = batch["batch_kwargs"]
+                    for suite_name in batch["expectation_suite_names"]:
+                        suite = context.get_expectation_suite(suite_name)
+                        batch = context.get_batch(batch_kwargs, suite)
+                        assets_to_validate.append(batch)
+            else:
+                assets_to_validate.append(
+                    context.get_batch(batch_kwargs, expectation_suite_name)
+                )
 
+        # Run validation operator
         results = context.run_validation_operator(
-            checkpoint["validation_operator_name"],
-            assets_to_validate=batches_to_validate,
-            run_id={"run_name": prefect.context.get("task_slug")},
+            checkpoint_name or ge_checkpoint["validation_operator_name"],
+            assets_to_validate=assets_to_validate,
+            run_id={"run_name": run_name or prefect.context.get("task_slug")},
         )
 
         if results.success is False:
-            raise signals.VALIDATIONFAIL(result=results)
+            raise signals.FAIL(result=results)
+
+        # Generate artifact markdown
+        if not disable_markdown_artifact:
+            run_info_at_end = True
+            validation_results_page_renderer = (
+                ge.render.renderer.ValidationResultsPageRenderer(
+                    run_info_at_end=run_info_at_end
+                )
+            )
+            rendered_document_content_list = (
+                validation_results_page_renderer.render_validation_operator_result(
+                    validation_operator_result=results
+                )
+            )
+            markdown_artifact = " ".join(
+                ge.render.view.DefaultMarkdownPageView().render(
+                    rendered_document_content_list
+                )
+            )
+
+            create_markdown(markdown_artifact)
 
         return results


Test patch:
diff --git a/tests/tasks/great_expectations/test_great_expectations_checkpoints.py b/tests/tasks/great_expectations/test_great_expectations_checkpoints.py
--- a/tests/tasks/great_expectations/test_great_expectations_checkpoints.py
+++ b/tests/tasks/great_expectations/test_great_expectations_checkpoints.py
@@ -17,18 +17,32 @@ def test_inits_with_no_args(self):
     def test_kwargs_get_passed_to_task_init(self):
         t = RunGreatExpectationsCheckpoint(
             checkpoint_name="checkpoint",
+            context=1234,
+            assets_to_validate=["assets"],
+            batch_kwargs={"kwargs": "here"},
+            expectation_suite_name="name",
+            get_checkpoint_from_context=True,
             context_root_dir="/path/to/somewhere",
             runtime_environment={
                 "plugins_directory": "/path/to/plugins/somewhere/else"
             },
-            run_name="woooo",
+            run_name="1234",
+            run_info_at_end=False,
+            disable_markdown_artifact=True,
         )
         assert t.checkpoint_name == "checkpoint"
+        assert t.context == 1234
+        assert t.assets_to_validate == ["assets"]
+        assert t.batch_kwargs == {"kwargs": "here"}
+        assert t.expectation_suite_name == "name"
+        assert t.get_checkpoint_from_context == True
         assert t.context_root_dir == "/path/to/somewhere"
         assert t.runtime_environment == {
             "plugins_directory": "/path/to/plugins/somewhere/else"
         }
-        assert t.run_name == "woooo"
+        assert t.run_name == "1234"
+        assert t.run_info_at_end == False
+        assert t.disable_markdown_artifact == True
 
     def test_raises_if_checkpoint_not_provided(self, monkeypatch):
         task = RunGreatExpectationsCheckpoint()
@@ -37,42 +51,3 @@ def test_raises_if_checkpoint_not_provided(self, monkeypatch):
         monkeypatch.setattr("prefect.tasks.great_expectations", great_expectations)
         with pytest.raises(ValueError, match="checkpoint"):
             task.run()
-
-
-@pytest.fixture(scope="session")
-def example_context_root_dir():
-    """
-    Copies the examples configuration directory to the proper location for the test and exposes its path for the test.
-    """
-    with tempfile.TemporaryDirectory() as tmp:
-        examples_path = os.path.abspath(
-            "examples/task_library/great_expectations/great_expectations/"
-        )
-        shutil.copytree(examples_path, tmp, dirs_exist_ok=True)
-        yield tmp
-
-
-@pytest.mark.skipif("sys.version_info <= (3,8)")
-class TestIntegrationWithLocalConfig:
-    """
-    Integration test to check the task library task still integrates with how Great Expectations checkpoints are configured.
-    This only runs on 3.8 since it uses a shutil kwarg (dirs_exist_ok) only available in 3.8+ for the purposes of the test.
-    """
-
-    def test_runs_with_different_context_root_dir(self, example_context_root_dir):
-        t = RunGreatExpectationsCheckpoint(
-            "warning_npi_checkpoint_circle", context_root_dir=example_context_root_dir
-        )
-
-        result = t.run()
-
-        assert result["success"]
-
-    def test_raises_if_validation_fails(self, example_context_root_dir):
-        t = RunGreatExpectationsCheckpoint(
-            "guaranteed_failure_npi_checkpoint_circle",
-            context_root_dir=example_context_root_dir,
-        )
-
-        with pytest.raises(prefect.engine.signals.VALIDATIONFAIL):
-            result = t.run()