Repository: great-expectations/great_expectations
Pull request #6826 (2023-01-18 14:44:14)

[MAINTENANCE] Standardize Store CRUD
Changes proposed in this pull request:
- Ensure that `add_store` always persists
- Add `delete_store`


### 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.


codesee-maps[bot]: <details>
<summary><b>👇 <a href="https://app.codesee.io/r/reviews?pr=6826&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=6826&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/6826/99c37fc0/3f2e81860fb6bbb6756da38025f01f574e8ef90c.svg)](https://app.codesee.io/r/reviews?pr=6826&src=https%3A%2F%2Fgithub.com%2Fgreat-expectations%2Fgreat_expectations)

#### Legend
![CodeSee Map Legend](https://app.codesee.io/Legend_v2.png)
Kilo59: Github is being pretty uncool and not letting me comment on particular lines anymore.

Can you update the string `"Store"` return annotations to be non-strings with `from __future__ import annotations`?

Patch:
diff --git a/great_expectations/data_context/config_validator/yaml_config_validator.py b/great_expectations/data_context/config_validator/yaml_config_validator.py
index c4fb82b35..de4db68b7 100644
--- a/great_expectations/data_context/config_validator/yaml_config_validator.py
+++ b/great_expectations/data_context/config_validator/yaml_config_validator.py
@@ -381,12 +381,9 @@ class _YamlConfigValidator:
         """
         print(f"\tInstantiating as a Store, since class_name is {class_name}")
         store_name: str = name or config.get("name") or "my_temp_store"
-        instantiated_class = cast(
-            Store,
-            self._data_context._build_store_from_config(
-                store_name=store_name,
-                store_config=config,
-            ),
+        instantiated_class = self._data_context._build_store_from_config(
+            store_name=store_name,
+            store_config=config,
         )
         store_name = instantiated_class.store_name or store_name
         self._data_context.config["stores"][store_name] = config
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 3b8201a4a..2cfc66dee 100644
--- a/great_expectations/data_context/data_context/abstract_data_context.py
+++ b/great_expectations/data_context/data_context/abstract_data_context.py
@@ -99,7 +99,6 @@ from great_expectations.data_context.types.resource_identifiers import (
 )
 from great_expectations.data_context.util import (
     PasswordMasker,
-    build_store_from_config,
     instantiate_class_from_config,
     parse_substitution_variable,
 )
@@ -228,7 +227,7 @@ class AbstractDataContext(ConfigPeer, ABC):
         )
         # Init stores
         self._stores: dict = {}
-        self._init_stores(self.project_config_with_variables_substituted.stores)  # type: ignore[arg-type]
+        self._init_stores(self.project_config_with_variables_substituted.stores)
 
         # Init data_context_id
         self._data_context_id = self._construct_data_context_id()
@@ -567,7 +566,7 @@ class AbstractDataContext(ConfigPeer, ABC):
                     f"{float(CURRENT_GX_CONFIG_VERSION)} in order to use the new 'Profiler Store' feature.\n  "
                     f"Visit {AbstractDataContext.MIGRATION_WEBSITE} to learn more about the upgrade process."
                 )
-                built_store: Optional[Store] = self._build_store_from_config(
+                built_store: Store = self._build_store_from_config(
                     profiler_store_name,  # type: ignore[arg-type]
                     DataContextConfigDefaults.DEFAULT_STORES.value[profiler_store_name],  # type: ignore[index,arg-type]
                 )
@@ -1189,19 +1188,44 @@ class AbstractDataContext(ConfigPeer, ABC):
         masked_config: dict = PasswordMasker.sanitize_config(substituted_config)
         return masked_config
 
-    def add_store(self, store_name: str, store_config: dict) -> Optional[Store]:
+    def add_store(self, store_name: str, store_config: dict) -> Store:
         """Add a new Store to the DataContext and (for convenience) return the instantiated Store object.
 
         Args:
-            store_name (str): a key for the new Store in in self._stores
+            store_name (str): a key for the new Store in in self.stores
             store_config (dict): a config for the Store to add
 
         Returns:
             store (Store)
         """
+        store = self._build_store_from_config(store_name, store_config)
 
-        self.config.stores[store_name] = store_config  # type: ignore[index]
-        return self._build_store_from_config(store_name, store_config)
+        # Both the config and the actual stores need to be kept in sync
+        self.config.stores[store_name] = store_config
+        self._stores[store_name] = store
+
+        self._save_project_config()
+        return store
+
+    def delete_store(self, store_name: str) -> None:
+        """Delete an existing Store from the DataContext.
+
+        Args:
+            store_name: The name of the Store to be deleted.
+
+        Raises:
+            StoreConfigurationError if the target Store is not found.
+        """
+        if store_name not in self.config.stores and store_name not in self._stores:
+            raise gx_exceptions.StoreConfigurationError(
+                f'Attempted to delete a store named: "{store_name}". It is not a configured store.'
+            )
+
+        # Both the config and the actual stores need to be kept in sync
+        self.config.stores.pop(store_name, None)
+        self._stores.pop(store_name, None)
+
+        self._save_project_config()
 
     def list_datasources(self) -> List[dict]:
         """List currently-configured datasources on this context. Masks passwords.
@@ -2980,9 +3004,7 @@ Generated, evaluated, and stored {total_expectations} Expectations during profil
             conf_file_option="usage_statistics_url",
         )
 
-    def _build_store_from_config(
-        self, store_name: str, store_config: dict
-    ) -> Optional[Store]:
+    def _build_store_from_config(self, store_name: str, store_config: dict) -> Store:
         module_name = "great_expectations.data_context.store"
         # Set expectations_store.store_backend_id to the data_context_id from the project_config if
         # the expectations_store does not yet exist by:
@@ -3004,7 +3026,7 @@ Generated, evaluated, and stored {total_expectations} Expectations during profil
         ):
             store_config["store_backend"].update({"suppress_store_backend_id": True})
 
-        new_store = build_store_from_config(
+        new_store = Store.build_store_from_config(
             store_name=store_name,
             store_config=store_config,
             module_name=module_name,
@@ -3284,7 +3306,7 @@ Generated, evaluated, and stored {total_expectations} Expectations during profil
 
     def _construct_data_context_id(self) -> str:
         # Choose the id of the currently-configured expectations store, if it is a persistent store
-        expectations_store = self._stores[self.variables.expectations_store_name]
+        expectations_store = self.stores[self.expectations_store_name]
         if isinstance(expectations_store.store_backend, TupleStoreBackend):
             # suppress_warnings since a warning will already have been issued during the store creation
             # if there was an invalid store config
diff --git a/great_expectations/data_context/data_context/file_data_context.py b/great_expectations/data_context/data_context/file_data_context.py
index c2589b845..c73937583 100644
--- a/great_expectations/data_context/data_context/file_data_context.py
+++ b/great_expectations/data_context/data_context/file_data_context.py
@@ -15,7 +15,6 @@ from great_expectations.data_context.data_context_variables import (
     DataContextVariableSchema,
     FileDataContextVariables,
 )
-from great_expectations.data_context.store.store import Store
 from great_expectations.data_context.types.base import (
     DataContextConfig,
     datasourceConfigSchema,
@@ -135,15 +134,6 @@ class FileDataContext(SerializableDataContext):
         )
         return variables
 
-    def add_store(self, store_name: str, store_config: dict) -> Optional[Store]:
-        """
-        See parent `AbstractDataContext.add_store()` for more information.
-
-        """
-        store = super().add_store(store_name=store_name, store_config=store_config)
-        self._save_project_config()
-        return store
-
     @classmethod
     def _load_file_backed_project_config(
         cls,
diff --git a/great_expectations/data_context/store/store.py b/great_expectations/data_context/store/store.py
index 1347158be..a9259438d 100644
--- a/great_expectations/data_context/store/store.py
+++ b/great_expectations/data_context/store/store.py
@@ -1,6 +1,9 @@
+from __future__ import annotations
+
 import logging
 from typing import Any, Dict, List, Optional, Tuple, Type
 
+import great_expectations.exceptions as gx_exceptions
 from great_expectations.core.configuration import AbstractConfig
 from great_expectations.core.data_context_key import DataContextKey
 from great_expectations.data_context.store.gx_cloud_store_backend import (
@@ -212,3 +215,40 @@ class Store:
             name = config.name
 
         return self.store_backend.build_key(name=name, id=id)
+
+    @staticmethod
+    def build_store_from_config(
+        store_name: Optional[str] = None,
+        store_config: Optional[dict] = None,
+        module_name: str = "great_expectations.data_context.store",
+        runtime_environment: Optional[dict] = None,
+    ) -> Store:
+        if store_config is None or module_name is None:
+            raise gx_exceptions.StoreConfigurationError(
+                "Cannot build a store without both a store_config and a module_name"
+            )
+
+        try:
+            config_defaults: dict = {
+                "store_name": store_name,
+                "module_name": module_name,
+            }
+            new_store = instantiate_class_from_config(
+                config=store_config,
+                runtime_environment=runtime_environment,
+                config_defaults=config_defaults,
+            )
+        except gx_exceptions.DataContextError as e:
+            new_store = None
+            logger.critical(
+                f"Error {e} occurred while attempting to instantiate a store."
+            )
+        if not new_store:
+            class_name: str = store_config["class_name"]
+            module_name = store_config["module_name"]
+            raise gx_exceptions.ClassInstantiationError(
+                module_name=module_name,
+                package_name=None,
+                class_name=class_name,
+            )
+        return new_store
diff --git a/great_expectations/data_context/types/base.py b/great_expectations/data_context/types/base.py
index d5798f3a0..47ff29ad1 100644
--- a/great_expectations/data_context/types/base.py
+++ b/great_expectations/data_context/types/base.py
@@ -2387,7 +2387,7 @@ class DataContextConfig(BaseYamlConfig):
         self.plugins_directory = plugins_directory
         if validation_operators is not None:
             self.validation_operators = validation_operators
-        self.stores = stores
+        self.stores = stores or {}
         self.notebooks = notebooks
         self.data_docs_sites = data_docs_sites
         self.config_variables_file_path = config_variables_file_path
diff --git a/great_expectations/data_context/util.py b/great_expectations/data_context/util.py
index d9ace603a..1c240cfac 100644
--- a/great_expectations/data_context/util.py
+++ b/great_expectations/data_context/util.py
@@ -3,12 +3,11 @@ import inspect
 import logging
 import os
 import warnings
-from typing import TYPE_CHECKING, Any, Optional
+from typing import Any, Optional
 from urllib.parse import urlparse
 
 import pyparsing as pp
 
-import great_expectations.exceptions as gx_exceptions
 from great_expectations.types import safe_deep_copy
 from great_expectations.util import load_class, verify_dynamic_loading_support
 
@@ -17,9 +16,6 @@ try:
 except ImportError:
     sa = None
 
-if TYPE_CHECKING:
-    from great_expectations.data_context.store import Store
-
 logger = logging.getLogger(__name__)
 
 
@@ -104,39 +100,6 @@ def instantiate_class_from_config(config, runtime_environment, config_defaults=N
     return class_instance
 
 
-def build_store_from_config(
-    store_name: Optional[str] = None,
-    store_config: Optional[dict] = None,
-    module_name: str = "great_expectations.data_context.store",
-    runtime_environment: Optional[dict] = None,
-) -> Optional["Store"]:
-    if store_config is None or module_name is None:
-        return None
-
-    try:
-        config_defaults: dict = {
-            "store_name": store_name,
-            "module_name": module_name,
-        }
-        new_store = instantiate_class_from_config(
-            config=store_config,
-            runtime_environment=runtime_environment,
-            config_defaults=config_defaults,
-        )
-    except gx_exceptions.DataContextError as e:
-        new_store = None
-        logger.critical(f"Error {e} occurred while attempting to instantiate a store.")
-    if not new_store:
-        class_name: str = store_config["class_name"]
-        module_name = store_config["module_name"]
-        raise gx_exceptions.ClassInstantiationError(
-            module_name=module_name,
-            package_name=None,
-            class_name=class_name,
-        )
-    return new_store
-
-
 def format_dict_for_error_message(dict_):
     # TODO : Tidy this up a bit. Indentation isn't fully consistent.
 


Test patch:
diff --git a/tests/data_context/migrator/conftest.py b/tests/data_context/migrator/conftest.py
index 1347f3a82..783af4cb1 100644
--- a/tests/data_context/migrator/conftest.py
+++ b/tests/data_context/migrator/conftest.py
@@ -200,7 +200,7 @@ def empty_serialized_configuration_bundle() -> dict:
             },
             "notebooks": None,
             "plugins_directory": None,
-            "stores": None,
+            "stores": {},
             "validations_store_name": None,
         },
         "datasources": [],
@@ -235,7 +235,7 @@ def serialized_configuration_bundle() -> dict:
             },
             "notebooks": None,
             "plugins_directory": None,
-            "stores": None,
+            "stores": {},
             "validations_store_name": None,
         },
         "datasources": [
diff --git a/tests/data_context/store/test_store.py b/tests/data_context/store/test_store.py
index eeb346e49..d00732ca2 100644
--- a/tests/data_context/store/test_store.py
+++ b/tests/data_context/store/test_store.py
@@ -1,5 +1,6 @@
 import pytest
 
+import great_expectations.exceptions as gx_exceptions
 from great_expectations.core.configuration import AbstractConfig
 from great_expectations.data_context.store.store import Store
 
@@ -29,3 +30,42 @@ def test_store_deserialize() -> None:
     store = Store()
     value = {"a": "b"}
     assert store.deserialize(value) == value
+
+
+@pytest.mark.unit
+def test_build_store_from_config_success():
+    store_name = "my_new_store"
+    store_config = {
+        "module_name": "great_expectations.data_context.store",
+        "class_name": "ExpectationsStore",
+    }
+    store = Store.build_store_from_config(
+        store_name=store_name,
+        store_config=store_config,
+    )
+    assert isinstance(store, Store)
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "store_config,module_name",
+    [
+        pytest.param(None, "great_expectations.data_context.store", id="config_none"),
+        pytest.param(
+            {
+                "module_name": "great_expectations.data_context.store",
+                "class_name": "ExpectationsStore",
+            },
+            None,
+            id="module_name_none",
+        ),
+        pytest.param(None, None, id="config_and_module_name_both_none"),
+    ],
+)
+def test_build_store_from_config_failure(store_config: dict, module_name: str):
+    with pytest.raises(gx_exceptions.StoreConfigurationError):
+        Store.build_store_from_config(
+            store_name="my_new_store",
+            store_config=store_config,
+            module_name=module_name,
+        )
diff --git a/tests/data_context/test_data_context_state_management.py b/tests/data_context/test_data_context_state_management.py
new file mode 100644
index 000000000..79dbe0417
--- /dev/null
+++ b/tests/data_context/test_data_context_state_management.py
@@ -0,0 +1,94 @@
+import pytest
+
+from great_expectations.data_context.data_context.ephemeral_data_context import (
+    EphemeralDataContext,
+)
+from great_expectations.data_context.types.base import (
+    DataContextConfig,
+    InMemoryStoreBackendDefaults,
+)
+from great_expectations.exceptions.exceptions import StoreConfigurationError
+
+
+class EphemeralDataContextSpy(EphemeralDataContext):
+    """
+    Simply wraps around EphemeralDataContext but keeps tabs on specific method calls around state management.
+    """
+
+    def __init__(
+        self,
+        project_config: DataContextConfig,
+    ) -> None:
+        super().__init__(project_config)
+        self.save_count = 0
+
+    def _save_project_config(self):
+        """
+        No-op our persistence mechanism but increment an internal counter to ensure it was used.
+        """
+        self.save_count += 1
+
+
+@pytest.fixture
+def in_memory_data_context() -> EphemeralDataContextSpy:
+    config = DataContextConfig(store_backend_defaults=InMemoryStoreBackendDefaults())
+    context = EphemeralDataContextSpy(project_config=config)
+    return context
+
+
+@pytest.mark.unit
+def test_add_store(in_memory_data_context: EphemeralDataContextSpy):
+    context = in_memory_data_context
+
+    num_stores_before = len(context.stores)
+    num_store_configs_before = len(context.config.stores)
+
+    context.add_store(
+        store_name="my_new_store",
+        store_config={
+            "module_name": "great_expectations.data_context.store",
+            "class_name": "ExpectationsStore",
+        },
+    )
+
+    num_stores_after = len(context.stores)
+    num_store_configs_after = len(context.config.stores)
+
+    assert num_stores_after == num_stores_before + 1
+    assert num_store_configs_after == num_store_configs_before + 1
+    assert context.save_count == 1
+
+
+@pytest.mark.unit
+def test_delete_store_success(in_memory_data_context: EphemeralDataContextSpy):
+    context = in_memory_data_context
+
+    num_stores_before = len(context.stores)
+    num_store_configs_before = len(context.config.stores)
+
+    context.delete_store("checkpoint_store")  # We know this to be a default name
+
+    num_stores_after = len(context.stores)
+    num_store_configs_after = len(context.config.stores)
+
+    assert num_stores_after == num_stores_before - 1
+    assert num_store_configs_after == num_store_configs_before - 1
+    assert context.save_count == 1
+
+
+@pytest.mark.unit
+def test_delete_store_failure(in_memory_data_context: EphemeralDataContextSpy):
+    context = in_memory_data_context
+
+    num_stores_before = len(context.stores)
+    num_store_configs_before = len(context.config.stores)
+
+    with pytest.raises(StoreConfigurationError):
+        context.delete_store("my_fake_store_name")
+
+    num_stores_after = len(context.stores)
+    num_store_configs_after = len(context.config.stores)
+
+    assert num_stores_after == num_stores_before
+    assert num_store_configs_after == num_store_configs_before
+    assert context.save_count == 0
diff --git a/tests/data_context/test_data_context_utils.py b/tests/data_context/test_data_context_utils.py
index 0437a3573..40591e7fd 100644
--- a/tests/data_context/test_data_context_utils.py
+++ b/tests/data_context/test_data_context_utils.py
@@ -2,7 +2,7 @@ import os
 
 import pytest
 
-import great_expectations.exceptions as gee
+import great_expectations.exceptions as gx_exceptions
 from great_expectations.data_context.util import (
     PasswordMasker,
     parse_substitution_variable,
@@ -12,12 +12,12 @@ from great_expectations.util import load_class
 
 
 def test_load_class_raises_error_when_module_not_found():
-    with pytest.raises(gee.PluginModuleNotFoundError):
+    with pytest.raises(gx_exceptions.PluginModuleNotFoundError):
         load_class("foo", "bar")
 
 
 def test_load_class_raises_error_when_class_not_found():
-    with pytest.raises(gee.PluginClassNotFoundError):
+    with pytest.raises(gx_exceptions.PluginClassNotFoundError):
         load_class("TotallyNotARealClass", "great_expectations.datasource")
 
 
diff --git a/tests/test_utils.py b/tests/test_utils.py
index bf0137f6d..ac7d318b4 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -18,16 +18,14 @@ from great_expectations.data_context.store import (
     CheckpointStore,
     ConfigurationStore,
     ProfilerStore,
+    Store,
     StoreBackend,
 )
 from great_expectations.data_context.types.base import BaseYamlConfig, CheckpointConfig
 from great_expectations.data_context.types.resource_identifiers import (
     ConfigurationIdentifier,
 )
-from great_expectations.data_context.util import (
-    build_store_from_config,
-    instantiate_class_from_config,
-)
+from great_expectations.data_context.util import instantiate_class_from_config
 from great_expectations.execution_engine import SqlAlchemyExecutionEngine
 
 logger = logging.getLogger(__name__)
@@ -193,7 +191,7 @@ def build_tuple_filesystem_store_backend(
         "base_directory": base_directory,
     }
     store_backend_config.update(**kwargs)
-    return build_store_from_config(
+    return Store.build_store_from_config(
         store_config=store_backend_config,
         module_name=module_name,
         runtime_environment=None,
@@ -321,7 +319,7 @@ def build_configuration_store(
         "overwrite_existing": overwrite_existing,
         "store_backend": store_backend,
     }
-    configuration_store: ConfigurationStore = build_store_from_config(  # type: ignore[assignment]
+    configuration_store: ConfigurationStore = Store.build_store_from_config(  # type: ignore[assignment]
         store_config=store_config,
         module_name=module_name,
         runtime_environment=None,