Merge pull request #16 from great-expectations/m/CORE-751/configure_file_data_context_cleanup

[FEATURE] Allow configure_file_data_context to return a generator

Author: tyler-hoffman

great_expectations_provider/operators/validate_checkpoint.py

@@ -1,15 +1,16 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Callable, Literal
+import inspect
+from typing import TYPE_CHECKING, Callable, Generator, Literal, cast
 
 from airflow.models import BaseOperator
+from great_expectations.data_context import AbstractDataContext, FileDataContext
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
     from great_expectations import Checkpoint
     from great_expectations.checkpoint.checkpoint import CheckpointDescriptionDict
     from great_expectations.core.batch import BatchParameters
-    from great_expectations.data_context import AbstractDataContext, FileDataContext
 
 
 class GXValidateCheckpointOperator(BaseOperator):
@@ -18,7 +19,9 @@ def __init__(
         configure_checkpoint: Callable[[AbstractDataContext], Checkpoint],
         batch_parameters: BatchParameters | None = None,
         context_type: Literal["ephemeral", "cloud", "file"] = "ephemeral",
-        configure_file_data_context: Callable[[], FileDataContext] | None = None,
+        configure_file_data_context: Callable[[], FileDataContext]
+        | Callable[[], Generator[FileDataContext, None, None]]
+        | None = None,
         *args,
         **kwargs,
     ) -> None:
@@ -40,15 +43,50 @@ def execute(self, context: Context) -> CheckpointDescriptionDict:
         import great_expectations as gx
 
         gx_context: AbstractDataContext
+        file_context_generator: Generator[FileDataContext, None, None] | None = None
 
         if self.context_type == "file":
             if not self.configure_file_data_context:
                 raise ValueError(
                     "Parameter `configure_file_data_context` must be specified if `context_type` is `file`"
                 )
-            gx_context = self.configure_file_data_context()
+            elif inspect.isgeneratorfunction(self.configure_file_data_context):
+                file_context_generator = self.configure_file_data_context()
+                gx_context = self._get_value_from_generator(file_context_generator)
+            else:
+                file_context_fn = cast(
+                    Callable[[], FileDataContext], self.configure_file_data_context
+                )
+                gx_context = file_context_fn()
         else:
             gx_context = gx.get_context(mode=self.context_type)
         checkpoint = self.configure_checkpoint(gx_context)
         result = checkpoint.run(batch_parameters=self.batch_parameters)
+
+        if file_context_generator:
+            self._allow_generator_teardown(file_context_generator)
+
         return result.describe_dict()
+
+    def _get_value_from_generator(
+        self, generator: Generator[FileDataContext, None, None]
+    ) -> FileDataContext:
+        try:
+            return next(generator)
+        except StopIteration:
+            raise RuntimeError("Generator must yield exactly once; did not yield")
+
+    def _allow_generator_teardown(self, generator: Generator) -> None:
+        """Run the generator to completion to allow for any cleanup/teardown.
+
+        Also does some error handling to ensure the generator doesn't yield more than once.
+        """
+        try:
+            # Check if we have another yield (this is an error case)
+            next(generator)
+        except StopIteration:
+            pass
+        else:
+            raise RuntimeError(
+                "Generator must yield exactly once; yielded more than once"
+            )

tests/unit/test_validate_checkpoint_operator.py

@@ -1,11 +1,11 @@
 import json
-from typing import Literal
+from typing import Generator, Literal
 from unittest.mock import Mock
 
 import pandas as pd
 import pytest
 from great_expectations import Checkpoint, ExpectationSuite, ValidationDefinition
-from great_expectations.data_context import AbstractDataContext
+from great_expectations.data_context import AbstractDataContext, FileDataContext
 from great_expectations.expectations import ExpectColumnValuesToBeInSet
 from pytest_mock import MockerFixture
 
@@ -191,3 +191,118 @@ def test_batch_parameters(self) -> None:
 
         # assert
         mock_checkpoint.run.assert_called_once_with(batch_parameters=batch_parameters)
+
+    def test_configure_file_data_context_with_without_generator(self) -> None:
+        """Expect that configure_file_data_context can just return a DataContext"""
+        # arrange
+        mock_context = Mock(spec=AbstractDataContext)
+        setup = Mock()
+        teardown = Mock()
+        configure_checkpoint = Mock()
+
+        def configure_file_data_context() -> FileDataContext:
+            setup()
+            return mock_context
+
+        validate_checkpoint = GXValidateCheckpointOperator(
+            task_id="validate_checkpoint_success",
+            configure_checkpoint=configure_checkpoint,
+            context_type="file",
+            configure_file_data_context=configure_file_data_context,
+        )
+
+        # act
+        validate_checkpoint.execute(context={})
+
+        # assert
+        setup.assert_called_once()
+        configure_checkpoint.assert_called_once_with(mock_context)
+        teardown.assert_not_called()
+
+    def test_configure_file_data_context_with_generator(self) -> None:
+        """Expect that configure_file_data_context can return a generator that yeidls a DataContext."""
+        # arrange
+        mock_context = Mock(spec=AbstractDataContext)
+        setup = Mock()
+        teardown = Mock()
+        configure_checkpoint = Mock()
+
+        def configure_file_data_context() -> Generator[FileDataContext, None, None]:
+            setup()
+            yield mock_context
+            teardown()
+
+        validate_checkpoint = GXValidateCheckpointOperator(
+            task_id="validate_checkpoint_success",
+            configure_checkpoint=configure_checkpoint,
+            context_type="file",
+            configure_file_data_context=configure_file_data_context,
+        )
+
+        # act
+        validate_checkpoint.execute(context={})
+
+        # assert
+        setup.assert_called_once()
+        configure_checkpoint.assert_called_once_with(mock_context)
+        teardown.assert_called_once()
+
+    def test_configure_file_data_context_with_generator_no_yield(self) -> None:
+        """Expect that configure_file_data_context errors if it does not yield a DataContext."""
+        # arrange
+        mock_context = Mock(spec=AbstractDataContext)
+        setup = Mock()
+        teardown = Mock()
+        configure_checkpoint = Mock()
+
+        def configure_file_data_context() -> Generator[FileDataContext, None, None]:
+            setup()
+            if False:
+                # Force this to be a generator for the test
+                yield mock_context
+            teardown()
+
+        validate_checkpoint = GXValidateCheckpointOperator(
+            task_id="validate_checkpoint_success",
+            configure_checkpoint=configure_checkpoint,
+            context_type="file",
+            configure_file_data_context=configure_file_data_context,
+        )
+
+        # act
+        with pytest.raises(RuntimeError, match="did not yield"):
+            validate_checkpoint.execute(context={})
+
+        # assert
+        setup.assert_called_once()
+        configure_checkpoint.assert_not_called()
+        teardown.assert_called_once()
+
+    def test_configure_file_data_context_with_generator_multiple_yields(self) -> None:
+        """Expect that configure_file_data_context errors if it yields multiple times."""
+        # arrange
+        mock_context = Mock(spec=AbstractDataContext)
+        setup = Mock()
+        teardown = Mock()
+        configure_checkpoint = Mock()
+
+        def configure_file_data_context() -> Generator[FileDataContext, None, None]:
+            setup()
+            yield mock_context
+            yield mock_context
+            teardown()
+
+        validate_checkpoint = GXValidateCheckpointOperator(
+            task_id="validate_checkpoint_success",
+            configure_checkpoint=configure_checkpoint,
+            context_type="file",
+            configure_file_data_context=configure_file_data_context,
+        )
+
+        # act
+        with pytest.raises(RuntimeError, match="yielded more than once"):
+            validate_checkpoint.execute(context={})
+
+        # assert
+        setup.assert_called_once()
+        teardown.assert_not_called()