Repository: great-expectations/great_expectations
Pull request #6892 (2023-01-23 15:22:54)

[MAINTENANCE] Standardize `Datasource` CRUD
Changes proposed in this pull request:
- Add `add_or_update_datasource`


### 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 | 88e7091382272c0bd28e99ace8c6629cec801d0e |
|<span aria-hidden="true">🔍</span> Latest deploy log | https://app.netlify.com/sites/niobium-lead-7998/deploys/63cea651617fc100083bc85b |
codesee-maps[bot]: <details>
<summary><b>👇 <a href="https://app.codesee.io/r/reviews?pr=6892&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=6892&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/6892/584ede0a/88e7091382272c0bd28e99ace8c6629cec801d0e.svg)](https://app.codesee.io/r/reviews?pr=6892&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 7a07e4c96..4862b5dca 100644
--- a/great_expectations/data_context/data_context/abstract_data_context.py
+++ b/great_expectations/data_context/data_context/abstract_data_context.py
@@ -700,8 +700,8 @@ class AbstractDataContext(ConfigPeer, ABC):
         name: str,
         initialize: bool = True,
         save_changes: Optional[bool] = None,
-        **kwargs: Optional[dict],
-    ) -> Optional[Union[LegacyDatasource, BaseDatasource]]:
+        **kwargs,
+    ) -> LegacyDatasource | BaseDatasource | None:
         """Add a new Datasource to the data context, with configuration provided as kwargs.
 
         --Documentation--
@@ -730,7 +730,7 @@ class AbstractDataContext(ConfigPeer, ABC):
     ) -> Optional[Union[LegacyDatasource, BaseDatasource]]:
         save_changes = self._determine_save_changes_flag(save_changes)
 
-        logger.debug(f"Starting BaseDataContext.add_datasource for {name}")
+        logger.debug(f"Starting AbstractDataContext.add_datasource for {name}")
 
         module_name: str = kwargs.get("module_name", "great_expectations.datasource")  # type: ignore[assignment]
         verify_dynamic_loading_support(module_name=module_name)
@@ -762,27 +762,85 @@ class AbstractDataContext(ConfigPeer, ABC):
     def update_datasource(
         self,
         datasource: Union[LegacyDatasource, BaseDatasource],
-        save_changes: Optional[bool] = None,
+        save_changes: bool | None = None,
     ) -> None:
         """
-        Updates a DatasourceConfig that already exists in the store.
+        Updates a Datasource that already exists in the store.
 
         Args:
-            datasource_config: The config object to persist using the DatasourceStore.
+            datasource: The Datasource object to update.
             save_changes: do I save changes to disk?
         """
         save_changes = self._determine_save_changes_flag(save_changes)
+        self._update_datasource(datasource=datasource, save_changes=save_changes)
 
-        datasource_config_dict: dict = datasourceConfigSchema.dump(datasource.config)
+    def _update_datasource(
+        self,
+        datasource: LegacyDatasource | BaseDatasource,
+        save_changes: bool | None = None,
+        **kwargs,
+    ) -> Datasource:
+        name = datasource.name
+        config = datasource.config
+        # `instantiate_class_from_config` requires `class_name`
+        config["class_name"] = datasource.__class__.__name__
+        config.update(kwargs)
+
+        datasource_config_dict: dict = datasourceConfigSchema.dump(config)
         datasource_config = DatasourceConfig(**datasource_config_dict)
-        datasource_name: str = datasource.name
 
         if save_changes:
             self._datasource_store.update_by_name(  # type: ignore[attr-defined]
-                datasource_name=datasource_name, datasource_config=datasource_config
+                datasource_name=name, datasource_config=datasource_config
+            )
+
+        updated_datasource = (
+            self._instantiate_datasource_from_config_and_update_project_config(
+                config=datasource_config, initialize=True, save_changes=False
             )
-        self.config.datasources[datasource_name] = datasource_config  # type: ignore[assignment,index]
-        self._cached_datasources[datasource_name] = datasource_config
+        )
+
+        # Invariant based on `initalize=True` above
+        assert updated_datasource is not None
+
+        return updated_datasource
+
+    def add_or_update_datasource(
+        self,
+        name: str,
+        **kwargs,
+    ) -> LegacyDatasource | BaseDatasource:
+        """
+        Add a new Datasource or update an existing one on the context depending on whether
+        it already exists or not. The configuration is provided as kwargs.
+
+        Args:
+            name: The name of the Datasource to add or update.
+            kwargs: Any relevant keyword args to use when adding or updating the target Datasource.
+
+        Returns:
+            The Datasource added or updated by the input `kwargs`.
+        """
+        existing_datasource = None
+        try:
+            existing_datasource = self.get_datasource(name)
+        except ValueError:
+            logger.info(
+                f"Could not find an existing datasource named '{name}'; creating a new one."
+            )
+
+        if existing_datasource:
+            result_datasource = self._update_datasource(
+                datasource=existing_datasource, **kwargs
+            )
+        else:
+            result_datasource = self.add_datasource(name=name, **kwargs)
+
+        # Invariant based on `initialize=True` in both add/update branches
+        assert result_datasource is not None
+
+        self._save_project_config()
+        return result_datasource
 
     def get_site_names(self) -> List[str]:
         """Get a list of configured site names."""
@@ -3595,7 +3653,7 @@ Generated, evaluated, and stored {total_expectations} Expectations during profil
         Raises:
             DatasourceInitializationError
         """
-        # Note that the call to `DatasourcStore.set` may alter the config object's state
+        # Note that the call to `DatasourceStore.set` may alter the config object's state
         # As such, we invoke it at the top of our function so any changes are reflected downstream
         if save_changes:
             config = self._datasource_store.set(key=None, value=config)  # type: ignore[attr-defined]
diff --git a/great_expectations/data_context/data_context/ephemeral_data_context.py b/great_expectations/data_context/data_context/ephemeral_data_context.py
index eab3ed61d..c510e6a40 100644
--- a/great_expectations/data_context/data_context/ephemeral_data_context.py
+++ b/great_expectations/data_context/data_context/ephemeral_data_context.py
@@ -1,5 +1,5 @@
 import logging
-from typing import Mapping, Optional, Union
+from typing import Dict, Mapping, Optional, Union, cast
 
 from great_expectations.core._docs_decorators import public_api
 from great_expectations.core.serializer import DictConfigSerializer
@@ -11,6 +11,7 @@ from great_expectations.data_context.data_context_variables import (
 )
 from great_expectations.data_context.types.base import (
     DataContextConfig,
+    DatasourceConfig,
     datasourceConfigSchema,
 )
 
@@ -65,4 +66,9 @@ class EphemeralDataContext(AbstractDataContext):
             store_backend=store_backend,
             serializer=DictConfigSerializer(schema=datasourceConfigSchema),
         )
+        # As the store is in-memory, it needs to be populated immediately
+        datasources = cast(Dict[str, DatasourceConfig], self.config.datasources or {})
+        for name, config in datasources.items():
+            datasource_store.set_by_name(datasource_name=name, datasource_config=config)
+
         self._datasource_store = datasource_store


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 32d36afab..3a6ce5ffb 100644
--- a/tests/data_context/test_data_context_state_management.py
+++ b/tests/data_context/test_data_context_state_management.py
@@ -9,6 +9,7 @@ from great_expectations.data_context.data_context.ephemeral_data_context import
 )
 from great_expectations.data_context.types.base import (
     DataContextConfig,
+    DatasourceConfig,
     InMemoryStoreBackendDefaults,
     ProgressBarsConfig,
 )
@@ -35,8 +36,21 @@ class EphemeralDataContextSpy(EphemeralDataContext):
 
 
 @pytest.fixture
-def in_memory_data_context() -> EphemeralDataContextSpy:
-    config = DataContextConfig(store_backend_defaults=InMemoryStoreBackendDefaults())
+def datasource_name() -> str:
+    return "my_pandas_datasource"
+
+
+@pytest.fixture
+def in_memory_data_context(
+    datasource_name: str,
+    datasource_config: DatasourceConfig,
+) -> EphemeralDataContextSpy:
+    datasources = {
+        datasource_name: datasource_config,
+    }
+    config = DataContextConfig(
+        datasources=datasources, store_backend_defaults=InMemoryStoreBackendDefaults()
+    )
     context = EphemeralDataContextSpy(project_config=config)
     return context
 
@@ -122,3 +136,76 @@ def test_update_project_config(
     context.update_project_config(config)
 
     assert context.progress_bars["globally"] is True
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "kwargs,expected_id",
+    [
+        pytest.param(
+            {},
+            None,
+            id="no kwargs",
+        ),
+        pytest.param(
+            {
+                "id": "d53c2384-f973-4a0c-9c85-af1d67c06f58",
+            },
+            "d53c2384-f973-4a0c-9c85-af1d67c06f58",
+            id="kwargs",
+        ),
+    ],
+)
+def test_add_or_update_datasource_updates_successfully(
+    in_memory_data_context: EphemeralDataContextSpy,
+    datasource_name: str,
+    kwargs: dict,
+    expected_id: str | None,
+):
+    context = in_memory_data_context
+
+    num_datasource_before = len(context.datasources)
+    num_datasource_configs_before = len(context.config.datasources)
+
+    assert (
+        datasource_name in context.datasources
+    ), f"Downstream logic in the test relies on {datasource_name} being a datasource; please check your fixtures."
+
+    datasource = context.add_or_update_datasource(name=datasource_name, **kwargs)
+    # Let's `id` as an example attr to change so we don't need to assert against the whole config
+    assert datasource.config["id"] == expected_id
+
+    num_datasource_after = len(context.datasources)
+    num_datasource_configs_after = len(context.config.datasources)
+
+    assert num_datasource_after == num_datasource_before
+    assert num_datasource_configs_after == num_datasource_configs_before
+    assert context.save_count == 1
+
+
+@pytest.mark.unit
+def test_add_or_update_datasource_adds_successfully(
+    in_memory_data_context: EphemeralDataContextSpy,
+):
+    context = in_memory_data_context
+
+    num_datasource_before = len(context.datasources)
+    num_datasource_configs_before = len(context.config.datasources)
+
+    datasource_name = "my_brand_new_datasource"
+
+    assert datasource_name not in context.datasources
+
+    _ = context.add_or_update_datasource(
+        name=datasource_name,
+        module_name="great_expectations.datasource",
+        class_name="PandasDatasource",
+    )
+
+    num_datasource_after = len(context.datasources)
+    num_datasource_configs_after = len(context.config.datasources)
+
+    assert datasource_name in context.datasources
+    assert num_datasource_after == num_datasource_before + 1
+    assert num_datasource_configs_after == num_datasource_configs_before + 1
+    assert context.save_count == 1