Repository: great-expectations/great_expectations
Pull request #7070 (2023-02-07 17:47:29)

[FEATURE] Add `datasource` arg to `DataContext` `Datasource` CRUD
Changes proposed in this pull request:
- Allow users to pass in an existing datasource through `add_datasource` and `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 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 | f4e813cb614e30903bdf9c3c21316e0382730ed4 |
|<span aria-hidden="true">🔍</span> Latest deploy log | https://app.netlify.com/sites/niobium-lead-7998/deploys/63e28eb31b5e8800089639d6 |
codesee-maps[bot]: <details>
<summary><b>👇 <a href="https://app.codesee.io/r/reviews?pr=7070&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=7070&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/7070/690bdd4e/f4e813cb614e30903bdf9c3c21316e0382730ed4.svg)](https://app.codesee.io/r/reviews?pr=7070&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/core/usage_statistics/usage_statistics.py b/great_expectations/core/usage_statistics/usage_statistics.py
index 093967618..66418e891 100644
--- a/great_expectations/core/usage_statistics/usage_statistics.py
+++ b/great_expectations/core/usage_statistics/usage_statistics.py
@@ -22,12 +22,6 @@ import requests
 from great_expectations import __version__ as ge_version
 from great_expectations.core import ExpectationSuite
 from great_expectations.core.usage_statistics.anonymizers.anonymizer import Anonymizer
-from great_expectations.core.usage_statistics.anonymizers.types.base import (
-    CLISuiteInteractiveFlagCombinations,  # noqa: TCH001
-)
-from great_expectations.core.usage_statistics.events import (
-    UsageStatsEvents,  # noqa: TCH001
-)
 from great_expectations.core.usage_statistics.execution_environment import (
     GXExecutionEnvironment,
     PackageInfo,
@@ -42,11 +36,20 @@ from great_expectations.rule_based_profiler.config import RuleBasedProfilerConfi
 
 if TYPE_CHECKING:
     from great_expectations.checkpoint.checkpoint import Checkpoint
+    from great_expectations.core.usage_statistics.anonymizers.types.base import (
+        CLISuiteInteractiveFlagCombinations,
+    )
+    from great_expectations.core.usage_statistics.events import (
+        UsageStatsEvents,
+    )
     from great_expectations.data_context import AbstractDataContext
+    from great_expectations.datasource import LegacyDatasource
+    from great_expectations.datasource.new_datasource import BaseDatasource
     from great_expectations.rule_based_profiler.rule_based_profiler import (
         RuleBasedProfiler,
     )
 
+
 STOP_SIGNAL = object()
 
 logger = logging.getLogger(__name__)
@@ -311,7 +314,8 @@ def usage_statistics_enabled_method(
             time_begin: int = int(round(time.time() * 1000))
             try:
                 if args_payload_fn is not None:
-                    nested_update(event_payload, args_payload_fn(*args, **kwargs))
+                    args_payload = args_payload_fn(*args, **kwargs) or {}
+                    nested_update(event_payload, args_payload)
 
                 result = func(*args, **kwargs)
                 message["success"] = True
@@ -449,7 +453,10 @@ def edit_expectation_suite_usage_statistics(
 
 
 def add_datasource_usage_statistics(
-    data_context: AbstractDataContext, name: str, **kwargs
+    data_context: AbstractDataContext,
+    name: str | None = None,
+    datasource: LegacyDatasource | BaseDatasource | None = None,
+    **kwargs,
 ) -> dict:
     if not data_context._usage_statistics_handler:
         return {}
@@ -470,6 +477,10 @@ def add_datasource_usage_statistics(
     payload = {}
     # noinspection PyBroadException
     try:
+        assert (
+            name or datasource
+        ), "Guaranteed to have either one of these values due to prior validation"
+        name = name or datasource.name
         payload = datasource_anonymizer._anonymize_datasource_info(name, kwargs)
     except Exception as e:
         logger.debug(
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 308f9e020..bd8c5cb17 100644
--- a/great_expectations/data_context/data_context/abstract_data_context.py
+++ b/great_expectations/data_context/data_context/abstract_data_context.py
@@ -718,17 +718,53 @@ class AbstractDataContext(ConfigPeer, ABC):
         self._cached_datasources[datasource_name] = updated_datasource
         return updated_datasource
 
+    @overload
+    def add_datasource(
+        self,
+        name: str = ...,
+        initialize: bool = ...,
+        save_changes: bool | None = ...,
+        datasource: None = ...,
+        **kwargs,
+    ) -> LegacyDatasource | BaseDatasource | None:
+        """
+        A `name` is provided.
+        `datasource` should not be provided.
+        """
+        ...
+
+    @overload
+    def add_datasource(
+        self,
+        name: None = ...,
+        initialize: bool = ...,
+        save_changes: bool | None = ...,
+        datasource: LegacyDatasource | BaseDatasource = ...,
+        **kwargs,
+    ) -> LegacyDatasource | BaseDatasource | None:
+        """
+        A `datasource` is provided.
+        `name` should not be provided.
+        """
+        ...
+
     @public_api
     @deprecated_argument(argument_name="save_changes", version="0.15.32")
+    @new_argument(
+        argument_name="datasource",
+        version="0.15.49",
+        message="Pass in an existing Datasource instead of individual constructor arguments",
+    )
     @usage_statistics_enabled_method(
         event_name=UsageStatsEvents.DATA_CONTEXT_ADD_DATASOURCE,
         args_payload_fn=add_datasource_usage_statistics,
     )
     def add_datasource(
         self,
-        name: str,
+        name: str | None = None,
         initialize: bool = True,
-        save_changes: Optional[bool] = None,
+        save_changes: bool | None = None,
+        datasource: LegacyDatasource | BaseDatasource | None = None,
         **kwargs,
     ) -> LegacyDatasource | BaseDatasource | None:
         """Add a new Datasource to the data context, with configuration provided as kwargs.
@@ -741,52 +777,75 @@ class AbstractDataContext(ConfigPeer, ABC):
             initialize: if False, add the Datasource to the config, but do not
                 initialize it, for example if a user needs to debug database connectivity.
             save_changes: should GX save the Datasource config?
+            datasource: an existing Datasource you wish to persist
             kwargs: the configuration for the new Datasource
 
         Returns:
             Datasource instance added.
         """
+        self._validate_add_datasource_args(name=name, datasource=datasource)
         return self._add_datasource(
-            name=name, initialize=initialize, save_changes=save_changes, **kwargs
+            name=name,
+            initialize=initialize,
+            save_changes=save_changes,
+            datasource=datasource,
+            **kwargs,
         )
 
+    @staticmethod
+    def _validate_add_datasource_args(
+        name: str | None, datasource: LegacyDatasource | BaseDatasource | None
+    ) -> None:
+        if not ((datasource is None) ^ (name is None)):
+            raise ValueError(
+                "Must either pass in an existing datasource or individual constructor arguments (but not both)"
+            )
+
     def _add_datasource(
         self,
-        name: str,
+        name: str | None = None,
         initialize: bool = True,
-        save_changes: Optional[bool] = None,
-        **kwargs: Optional[dict],
-    ) -> Optional[Union[LegacyDatasource, BaseDatasource]]:
+        save_changes: bool | None = None,
+        datasource: LegacyDatasource | BaseDatasource | None = None,
+        **kwargs,
+    ) -> LegacyDatasource | BaseDatasource | None:
         save_changes = self._determine_save_changes_flag(save_changes)
 
         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)
-        class_name: Optional[str] = kwargs.get("class_name")  # type: ignore[assignment]
-        datasource_class = load_class(module_name=module_name, class_name=class_name)  # type: ignore[arg-type]
-
-        # For any class that should be loaded, it may control its configuration construction
-        # by implementing a classmethod called build_configuration
-        config: Union[CommentedMap, dict]
-        if hasattr(datasource_class, "build_configuration"):
-            config = datasource_class.build_configuration(**kwargs)
+        if datasource:
+            config = datasource.config
         else:
-            config = kwargs
+            module_name: str = kwargs.get(
+                "module_name", "great_expectations.datasource"
+            )
+            verify_dynamic_loading_support(module_name=module_name)
+            class_name = kwargs.get("class_name", "Datasource")
+            datasource_class = load_class(
+                module_name=module_name, class_name=class_name
+            )
+
+            # For any class that should be loaded, it may control its configuration construction
+            # by implementing a classmethod called build_configuration
+            if hasattr(datasource_class, "build_configuration"):
+                config = datasource_class.build_configuration(**kwargs)
+            else:
+                config = kwargs
 
         datasource_config: DatasourceConfig = datasourceConfigSchema.load(
             CommentedMap(**config)
         )
-        datasource_config.name = name
+        if name:
+            datasource_config.name = name
 
-        datasource: Optional[
-            Union[LegacyDatasource, BaseDatasource]
-        ] = self._instantiate_datasource_from_config_and_update_project_config(
-            config=datasource_config,
-            initialize=initialize,
-            save_changes=save_changes,
+        return_datasource: LegacyDatasource | BaseDatasource | None = (
+            self._instantiate_datasource_from_config_and_update_project_config(
+                config=datasource_config,
+                initialize=initialize,
+                save_changes=save_changes,
+            )
         )
-        return datasource
+        return return_datasource
 
     @public_api
     def update_datasource(
@@ -809,14 +868,12 @@ class AbstractDataContext(ConfigPeer, ABC):
     def _update_datasource(
         self,
         datasource: LegacyDatasource | BaseDatasource,
-        save_changes: bool | None = None,
-        **kwargs,
+        save_changes: bool,
     ) -> 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)
@@ -828,7 +885,9 @@ class AbstractDataContext(ConfigPeer, ABC):
 
         updated_datasource = (
             self._instantiate_datasource_from_config_and_update_project_config(
-                config=datasource_config, initialize=True, save_changes=False
+                config=datasource_config,
+                initialize=True,
+                save_changes=False,
             )
         )
 
@@ -837,11 +896,38 @@ class AbstractDataContext(ConfigPeer, ABC):
 
         return updated_datasource
 
+    @overload
+    def add_or_update_datasource(
+        self,
+        name: str = ...,
+        datasource: None = ...,
+        **kwargs,
+    ) -> LegacyDatasource | BaseDatasource | None:
+        """
+        A `name` is provided.
+        `datasource` should not be provided.
+        """
+        ...
+
+    @overload
+    def add_or_update_datasource(
+        self,
+        name: None = ...,
+        datasource: LegacyDatasource | BaseDatasource = ...,
+        **kwargs,
+    ) -> LegacyDatasource | BaseDatasource | None:
+        """
+        A `datasource` is provided.
+        `name` should not be provided.
+        """
+        ...
+
     @public_api
     @new_method_or_class(version="0.15.48")
     def add_or_update_datasource(
         self,
-        name: str,
+        name: str | None = None,
+        datasource: LegacyDatasource | BaseDatasource | None = None,
         **kwargs,
     ) -> LegacyDatasource | BaseDatasource:
         """Add a new Datasource or update an existing one on the context depending on whether
@@ -849,35 +935,21 @@ class AbstractDataContext(ConfigPeer, ABC):
 
         Args:
             name: The name of the Datasource to add or update.
+            datasource: an existing Datasource you wish to persist.
             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: LegacyDatasource | BaseDatasource | XDatasource | None = (
-            None
+        self._validate_add_datasource_args(name=name, datasource=datasource)
+        datasource = self._add_datasource(
+            name=name,
+            datasource=datasource,
+            **kwargs,
         )
-        try:
-            existing_datasource = self.get_datasource(name)
-        except ValueError:
-            logger.info(
-                f"Could not find an existing datasource named '{name}'; creating a new one."
-            )
+        assert datasource is not None
 
-        if existing_datasource and isinstance(
-            existing_datasource, (LegacyDatasource, BaseDatasource)
-        ):
-            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
+        return datasource
 
     def get_site_names(self) -> List[str]:
         """Get a list of configured site names."""
@@ -4397,8 +4469,9 @@ Generated, evaluated, and stored {total_expectations} Expectations during profil
                 raw_config=raw_config, substituted_config=substituted_config
             )
         except Exception as e:
+            name = substituted_config.name or ""
             raise gx_exceptions.DatasourceInitializationError(
-                datasource_name=substituted_config.name, message=str(e)
+                datasource_name=name, message=str(e)
             )
         return datasource
 
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 c510e6a40..34a8fd297 100644
--- a/great_expectations/data_context/data_context/ephemeral_data_context.py
+++ b/great_expectations/data_context/data_context/ephemeral_data_context.py
@@ -69,6 +69,6 @@ class EphemeralDataContext(AbstractDataContext):
         # 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)
+            datasource_store.add_by_name(datasource_name=name, datasource_config=config)
 
         self._datasource_store = datasource_store
diff --git a/great_expectations/data_context/store/datasource_store.py b/great_expectations/data_context/store/datasource_store.py
index 1252449a5..045c34bc7 100644
--- a/great_expectations/data_context/store/datasource_store.py
+++ b/great_expectations/data_context/store/datasource_store.py
@@ -1,13 +1,13 @@
 from __future__ import annotations
 
 import copy
-from typing import List, Optional, Union
+from typing import TYPE_CHECKING, List, Optional, Union
 
+import great_expectations.exceptions as gx_exceptions
 from great_expectations.core.data_context_key import (
     DataContextKey,
     DataContextVariableKey,
 )
-from great_expectations.core.serializer import AbstractConfigSerializer  # noqa: TCH001
 from great_expectations.data_context.store.store import Store
 from great_expectations.data_context.store.store_backend import StoreBackend
 from great_expectations.data_context.types.base import (
@@ -15,11 +15,14 @@ from great_expectations.data_context.types.base import (
     datasourceConfigSchema,
 )
 from great_expectations.data_context.types.refs import GXCloudResourceRef
-from great_expectations.data_context.types.resource_identifiers import (
-    GXCloudIdentifier,  # noqa: TCH001
-)
 from great_expectations.util import filter_properties_dict
 
+if TYPE_CHECKING:
+    from great_expectations.core.serializer import AbstractConfigSerializer
+    from great_expectations.data_context.types.resource_identifiers import (
+        GXCloudIdentifier,
+    )
+
 
 class DatasourceStore(Store):
     """
@@ -144,22 +147,8 @@ class DatasourceStore(Store):
             id=datasource_config.id,
         )
 
-    def set_by_name(
-        self, datasource_name: str, datasource_config: DatasourceConfig
-    ) -> None:
-        """Persists a DatasourceConfig in the store by a given name.
-
-        Args:
-            datasource_name: The name of the Datasource to update.
-            datasource_config: The config object to persist using the StoreBackend.
-        """
-        datasource_key: DataContextVariableKey = self._determine_datasource_key(
-            datasource_name=datasource_name
-        )
-        self.set(datasource_key, datasource_config)
-
     def set(  # type: ignore[override]
-        self, key: Union[DataContextKey, None], value: DatasourceConfig, **_: dict
+        self, key: Union[DataContextKey, None], value: DatasourceConfig, **kwargs
     ) -> DatasourceConfig:
         """Create a datasource config in the store using a store_backend-specific key.
         Args:
@@ -171,10 +160,16 @@ class DatasourceStore(Store):
         """
         if not key:
             key = self._build_key_from_config(value)
+        return self._persist_datasource(key=key, config=value)
 
+    def _persist_datasource(
+        self, key: DataContextKey, config: DatasourceConfig
+    ) -> DatasourceConfig:
         # Make two separate requests to set and get in order to obtain any additional
         # values that may have been added to the config by the StoreBackend (i.e. object ids)
-        ref: Optional[Union[bool, GXCloudResourceRef]] = super().set(key, value)
+        ref: Optional[Union[bool, GXCloudResourceRef]] = super().set(
+            key=key, value=config
+        )
         if ref and isinstance(ref, GXCloudResourceRef):
             key.cloud_id = ref.cloud_id  # type: ignore[attr-defined]
 
@@ -186,6 +181,29 @@ class DatasourceStore(Store):
 
         return return_value
 
+    def add_by_name(
+        self, datasource_name: str, datasource_config: DatasourceConfig
+    ) -> None:
+        """Persists a DatasourceConfig in the store by a given name.
+
+        Args:
+            datasource_name: The name of the Datasource to update.
+            datasource_config: The config object to persist using the StoreBackend.
+
+        Raises:
+            DatasourceError: A DatasourceConfig with the given key already exists in the store.
+        """
+        datasource_key: DataContextVariableKey = self._determine_datasource_key(
+            datasource_name=datasource_name
+        )
+        try:
+            self.add(key=datasource_key, value=datasource_config)
+        except gx_exceptions.StoreBackendError:
+            raise gx_exceptions.DatasourceError(
+                datasource_name=datasource_name,
+                message="A Datasource with the given name already exists",
+            )
+
     def update_by_name(
         self, datasource_name: str, datasource_config: DatasourceConfig
     ) -> None:
@@ -196,20 +214,18 @@ class DatasourceStore(Store):
             datasource_config: The config object to persist using the StoreBackend.
 
         Raises:
-            ValueError if a DatasourceConfig is not found.
+            DatasourceNotFoundError: If a DatasourceConfig is not found.
         """
         datasource_key: DataContextVariableKey = self._determine_datasource_key(
             datasource_name=datasource_name
         )
-        if not self.has_key(datasource_key):  # noqa: W601
-            raise ValueError(
-                f"Unable to load datasource `{datasource_name}` -- no configuration found or invalid configuration."
+        try:
+            self.update(key=datasource_key, value=datasource_config)
+        except gx_exceptions.StoreBackendError:
+            raise gx_exceptions.DatasourceNotFoundError(
+                f"Could not find an existing Datasource named {datasource_name}."
             )
 
-        self.set_by_name(
-            datasource_name=datasource_name, datasource_config=datasource_config
-        )
-
     def _determine_datasource_key(self, datasource_name: str) -> DataContextVariableKey:
         datasource_key = DataContextVariableKey(
             resource_name=datasource_name,
diff --git a/great_expectations/exceptions/exceptions.py b/great_expectations/exceptions/exceptions.py
index 7076d619e..bc68aabc2 100644
--- a/great_expectations/exceptions/exceptions.py
+++ b/great_expectations/exceptions/exceptions.py
@@ -360,10 +360,9 @@ class BatchSpecError(DataContextError):
 
 
 class DatasourceError(DataContextError):
-    def __init__(self, datasource_name, message) -> None:
-        self.message = "Cannot initialize datasource {}, error: {}".format(
-            datasource_name,
-            message,
+    def __init__(self, datasource_name: str, message: str) -> None:
+        self.message = (
+            f"Cannot initialize datasource {datasource_name}, error: {message}"
         )
         super().__init__(self.message)
 
@@ -380,6 +379,10 @@ class DatasourceKeyPairAuthBadPassphraseError(DatasourceInitializationError):
     pass
 
 
+class DatasourceNotFoundError(DataContextError):
+    pass
+
+
 class InvalidConfigValueTypeError(DataContextError):
     pass
 


Test patch:
diff --git a/tests/core/usage_statistics/test_usage_statistics_data_context.py b/tests/core/usage_statistics/test_usage_statistics_data_context.py
index 943a7c380..fbe3567a4 100644
--- a/tests/core/usage_statistics/test_usage_statistics_data_context.py
+++ b/tests/core/usage_statistics/test_usage_statistics_data_context.py
@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import inspect
 import logging
 from unittest import mock
 
@@ -12,11 +11,15 @@ from great_expectations.core.usage_statistics.usage_statistics import ENABLED_ME
 logger = logging.getLogger(__name__)
 
 
-def filter_enabled_methods(
-    enabled_methods: dict[str, UsageStatsEvents]
-) -> dict[str, UsageStatsEvents]:
+@pytest.fixture
+def enable_usage_stats(monkeypatch):
+    monkeypatch.delenv("GE_USAGE_STATS")
+
+
+@pytest.mark.unit
+def test_enabled_methods_map_to_appropriate_usage_stats_events():
     filtered = {}
-    for qualified_method_name, event in enabled_methods.items():
+    for qualified_method_name, event in ENABLED_METHODS.items():
         parts = qualified_method_name.split(".")
         assert (
             len(parts) == 2
@@ -30,17 +33,7 @@ def filter_enabled_methods(
         method_name = parts[-1]
         filtered[method_name] = event
 
-    return filtered
-
-
-@pytest.fixture
-def enable_usage_stats(monkeypatch):
-    monkeypatch.delenv("GE_USAGE_STATS")
-
-
-@pytest.mark.unit
-def test_enabled_methods_map_to_appropriate_usage_stats_events():
-    assert filter_enabled_methods(ENABLED_METHODS) == {
+    assert filtered == {
         "add_datasource": UsageStatsEvents.DATA_CONTEXT_ADD_DATASOURCE,
         "build_data_docs": UsageStatsEvents.DATA_CONTEXT_BUILD_DATA_DOCS,
         "get_batch_list": UsageStatsEvents.DATA_CONTEXT_GET_BATCH_LIST,
@@ -57,8 +50,63 @@ def test_enabled_methods_map_to_appropriate_usage_stats_events():
     "great_expectations.core.usage_statistics.usage_statistics.UsageStatisticsHandler.emit"
 )
 @pytest.mark.parametrize(
-    "method_name,expected_event",
-    [(name, event) for name, event in filter_enabled_methods(ENABLED_METHODS).items()],
+    "method_name,expected_event,required_args",
+    [
+        pytest.param(
+            "add_datasource",
+            UsageStatsEvents.DATA_CONTEXT_ADD_DATASOURCE,
+            ["name"],
+            id="add datasource",
+        ),
+        pytest.param(
+            "build_data_docs",
+            UsageStatsEvents.DATA_CONTEXT_BUILD_DATA_DOCS,
+            [],
+            id="build_data_docs",
+        ),
+        pytest.param(
+            "get_batch_list",
+            UsageStatsEvents.DATA_CONTEXT_GET_BATCH_LIST,
+            [],
+            id="get_batch_list",
+        ),
+        pytest.param(
+            "open_data_docs",
+            UsageStatsEvents.DATA_CONTEXT_OPEN_DATA_DOCS,
+            [],
+            id="open_data_docs",
+        ),
+        pytest.param(
+            "run_checkpoint",
+            UsageStatsEvents.DATA_CONTEXT_RUN_CHECKPOINT,
+            [],
+            id="run_checkpoint",
+        ),
+        pytest.param(
+            "run_profiler_on_data",
+            UsageStatsEvents.DATA_CONTEXT_RUN_RULE_BASED_PROFILER_ON_DATA,
+            [],
+            id="run_profiler_on_data",
+        ),
+        pytest.param(
+            "run_profiler_with_dynamic_arguments",
+            UsageStatsEvents.DATA_CONTEXT_RUN_RULE_BASED_PROFILER_WITH_DYNAMIC_ARGUMENTS,
+            [],
+            id="run_profiler_with_dynamic_arguments",
+        ),
+        pytest.param(
+            "run_validation_operator",
+            UsageStatsEvents.DATA_CONTEXT_RUN_VALIDATION_OPERATOR,
+            ["validation_operator_name", "assets_to_validate"],
+            id="run_validation_operator",
+        ),
+        pytest.param(
+            "save_expectation_suite",
+            UsageStatsEvents.DATA_CONTEXT_SAVE_EXPECTATION_SUITE,
+            ["expectation_suite"],
+            id="save_expectation_suite",
+        ),
+    ],
 )
 @pytest.mark.parametrize(
     "data_context_fixture_name",
@@ -81,6 +129,7 @@ def test_all_relevant_context_methods_emit_usage_stats(
     enable_usage_stats,  # Needs to be before context fixtures to ensure usage stats handlers are attached
     method_name: str,
     expected_event: UsageStatsEvents,
+    required_args: list[str],
     data_context_fixture_name: str,
     request,
 ):
@@ -106,8 +155,7 @@ def test_all_relevant_context_methods_emit_usage_stats(
 
         # Generate a set of dummy values to trigger the target method
         # and invoke the usage stats decorator without causing side-effects
-        signature = inspect.signature(method)
-        kwargs = {param: None for param in signature.parameters}
+        kwargs = {param: mock.MagicMock() for param in required_args}
         method(**kwargs)
 
         mock_fn.assert_called_once()
diff --git a/tests/data_context/store/test_datasource_store.py b/tests/data_context/store/test_datasource_store.py
index e69ec1900..7e7514689 100644
--- a/tests/data_context/store/test_datasource_store.py
+++ b/tests/data_context/store/test_datasource_store.py
@@ -5,6 +5,7 @@ from unittest import mock
 
 import pytest
 
+import great_expectations.exceptions as gx_exceptions
 from great_expectations.core.data_context_key import DataContextVariableKey
 from great_expectations.core.serializer import (
     AbstractConfigSerializer,
@@ -279,14 +280,14 @@ def test_datasource_store_with_inline_store_backend(
 
 
 @pytest.mark.unit
-def test_datasource_store_set_by_name(
+def test_datasource_store_add_by_name(
     empty_datasource_store: DatasourceStore,
     datasource_config: DatasourceConfig,
     fake_datasource_name,
 ) -> None:
     assert len(empty_datasource_store.list_keys()) == 0
 
-    empty_datasource_store.set_by_name(
+    empty_datasource_store.add_by_name(
         datasource_name=fake_datasource_name, datasource_config=datasource_config
     )
 
@@ -395,13 +396,16 @@ def test_datasource_store_update_raises_error_if_datasource_doesnt_exist(
     empty_datasource_store: DatasourceStore,
 ) -> None:
     updated_datasource_config = DatasourceConfig()
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(gx_exceptions.DatasourceNotFoundError) as e:
         empty_datasource_store.update_by_name(
             datasource_name=fake_datasource_name,
             datasource_config=updated_datasource_config,
         )
 
-    assert f"Unable to load datasource `{fake_datasource_name}`" in str(e.value)
+    assert (
+        f"Could not find an existing Datasource named {fake_datasource_name}."
+        in str(e.value)
+    )
 
 
 @pytest.mark.integration
diff --git a/tests/data_context/test_data_context_datasources.py b/tests/data_context/test_data_context_datasources.py
index 2e7f1d182..dc760d663 100644
--- a/tests/data_context/test_data_context_datasources.py
+++ b/tests/data_context/test_data_context_datasources.py
@@ -50,7 +50,8 @@ def pandas_enabled_datasource_config() -> dict:
 
 
 @pytest.mark.integration
-def test_data_context_instantiates_ge_cloud_store_backend_with_cloud_config(
+@pytest.mark.cloud
+def test_data_context_instantiates_gx_cloud_store_backend_with_cloud_config(
     tmp_path: pathlib,
     data_context_config_with_datasources: DataContextConfig,
     ge_cloud_config: GXCloudConfig,
@@ -190,6 +191,7 @@ def test_get_datasource_cache_miss(
 
 
 @pytest.mark.unit
+@pytest.mark.cloud
 def test_DataContext_add_datasource_updates_cache_and_store(
     cloud_data_context_in_cloud_mode_with_datasource_pandas_engine: DataContext,
     datasource_config_with_names: DatasourceConfig,
@@ -222,6 +224,7 @@ def test_DataContext_add_datasource_updates_cache_and_store(
 
 
 @pytest.mark.unit
+@pytest.mark.cloud
 def test_DataContext_update_datasource_updates_existing_value_in_cache_and_store(
     cloud_data_context_in_cloud_mode_with_datasource_pandas_engine: DataContext,
     pandas_enabled_datasource_config: dict,
@@ -246,11 +249,11 @@ def test_DataContext_update_datasource_updates_existing_value_in_cache_and_store
     with mock.patch(
         "great_expectations.data_context.store.DatasourceStore.has_key"
     ), mock.patch(
-        "great_expectations.data_context.store.DatasourceStore.set"
-    ) as mock_set:
+        "great_expectations.data_context.store.DatasourceStore.update"
+    ) as mock_update:
         context.update_datasource(datasource)
 
-    mock_set.assert_called_once()
+    mock_update.assert_called_once()
     assert name in context.datasources
 
     with mock.patch(
@@ -263,6 +266,7 @@ def test_DataContext_update_datasource_updates_existing_value_in_cache_and_store
 
 
 @pytest.mark.unit
+@pytest.mark.cloud
 def test_DataContext_update_datasource_creates_new_value_in_cache_and_store(
     cloud_data_context_in_cloud_mode_with_datasource_pandas_engine: DataContext,
     pandas_enabled_datasource_config: dict,
@@ -285,15 +289,16 @@ def test_DataContext_update_datasource_creates_new_value_in_cache_and_store(
     with mock.patch(
         "great_expectations.data_context.store.DatasourceStore.has_key"
     ), mock.patch(
-        "great_expectations.data_context.store.DatasourceStore.set"
-    ) as mock_set:
+        "great_expectations.data_context.store.DatasourceStore.update"
+    ) as mock_update:
         context.update_datasource(datasource)
 
-    mock_set.assert_called_once()
+    mock_update.assert_called_once()
     assert name in context.datasources
 
 
 @pytest.mark.unit
+@pytest.mark.cloud
 def test_DataContext_delete_datasource_updates_cache(
     cloud_data_context_in_cloud_mode_with_datasource_pandas_engine: DataContext,
 ) -> None:
diff --git a/tests/data_context/test_data_context_state_management.py b/tests/data_context/test_data_context_state_management.py
index 4307c760c..fc0c55950 100644
--- a/tests/data_context/test_data_context_state_management.py
+++ b/tests/data_context/test_data_context_state_management.py
@@ -9,22 +9,37 @@ import great_expectations.exceptions as gx_exceptions
 from great_expectations.checkpoint.checkpoint import Checkpoint
 from great_expectations.core.expectation_configuration import ExpectationConfiguration
 from great_expectations.core.expectation_suite import ExpectationSuite
+from great_expectations.core.serializer import DictConfigSerializer
 from great_expectations.data_context.data_context.ephemeral_data_context import (
     EphemeralDataContext,
 )
 from great_expectations.data_context.store import ExpectationsStore, ProfilerStore
 from great_expectations.data_context.store.checkpoint_store import CheckpointStore
+from great_expectations.data_context.store.datasource_store import DatasourceStore
 from great_expectations.data_context.types.base import (
     CheckpointConfig,
     DataContextConfig,
     DatasourceConfig,
     InMemoryStoreBackendDefaults,
     ProgressBarsConfig,
+    datasourceConfigSchema,
 )
+from great_expectations.datasource.new_datasource import Datasource
 from great_expectations.exceptions.exceptions import StoreConfigurationError
 from great_expectations.rule_based_profiler.rule_based_profiler import RuleBasedProfiler
 
 
+class DatasourceStoreSpy(DatasourceStore):
+    def __init__(self) -> None:
+        self.save_count = 0
+        super().__init__(serializer=DictConfigSerializer(schema=datasourceConfigSchema))
+
+    def set(self, key, value, **kwargs):
+        ret = super().set(key=key, value=value, **kwargs)
+        self.save_count += 1
+        return ret
+
+
 class ExpectationsStoreSpy(ExpectationsStore):
     def __init__(self) -> None:
         self.save_count = 0
@@ -105,10 +120,15 @@ class EphemeralDataContextSpy(EphemeralDataContext):
     ) -> None:
         super().__init__(project_config)
         self.save_count = 0
+        self._datasource_store = DatasourceStoreSpy()
         self._expectations_store = ExpectationsStoreSpy()
         self._profiler_store = ProfilerStoreSpy()
         self._checkpoint_store = CheckpointStoreSpy()
 
+    @property
+    def datasource_store(self):
+        return self._datasource_store
+
     @property
     def expectations_store(self):
         return self._expectations_store
@@ -231,29 +251,68 @@ def test_update_project_config(
     assert context.progress_bars["globally"] is True
 
 
+@pytest.mark.unit
+def test_add_datasource_with_existing_datasource(
+    in_memory_data_context: EphemeralDataContextSpy,
+    datasource_config: DatasourceConfig,
+):
+    context = in_memory_data_context
+
+    config_dict = datasource_config.to_dict()
+    for attr in ("class_name", "module_name"):
+        config_dict.pop(attr)
+    config_dict["name"] = "my_datasource"
+
+    datasource = Datasource(**config_dict)
+    persisted_datasource = context.add_datasource(datasource=datasource)
+
+    expected_config = datasource.config
+    actual_config = persisted_datasource.config
+
+    # Name gets injected into data connector as part of serialization hook
+    # Removing for purposes of test assertion
+    data_connectors = actual_config["data_connectors"]
+    data_connector_name = tuple(data_connectors.keys())[0]
+    data_connectors[data_connector_name].pop("name")
+
+    assert actual_config == expected_config
+    assert context.datasource_store.save_count == 1
+
+
 @pytest.mark.unit
 @pytest.mark.parametrize(
-    "kwargs,expected_id",
+    "datasource,datasource_name",
     [
         pytest.param(
-            {},
-            None,
-            id="no kwargs",
-        ),
-        pytest.param(
-            {
-                "id": "d53c2384-f973-4a0c-9c85-af1d67c06f58",
-            },
-            "d53c2384-f973-4a0c-9c85-af1d67c06f58",
-            id="kwargs",
+            mock.MagicMock(),
+            "my_datasource",
+            id="both datasource and name",
         ),
+        pytest.param(None, None, id="neither datasource nor name"),
     ],
 )
-def test_add_or_update_datasource_updates_successfully(
+def test_add_datasource_conflicting_args_failure(
+    in_memory_data_context: EphemeralDataContextSpy,
+    datasource: mock.MagicMock | None,
+    datasource_name: str | None,
+):
+    context = in_memory_data_context
+
+    with pytest.raises(ValueError) as e:
+        context.add_datasource(datasource=datasource, name=datasource_name)
+
+    assert (
+        "an existing datasource or individual constructor arguments (but not both)"
+        in str(e.value)
+    )
+    assert context.datasource_store.save_count == 0
+
+
+@pytest.mark.unit
+def test_add_or_update_datasource_updates_with_individual_args_successfully(
     in_memory_data_context: EphemeralDataContextSpy,
     datasource_name: str,
-    kwargs: dict,
-    expected_id: str | None,
+    datasource_config: DatasourceConfig,
 ):
     context = in_memory_data_context
 
@@ -264,21 +323,72 @@ def test_add_or_update_datasource_updates_successfully(
         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)
+    config_dict = datasource_config.to_dict()
+    config_dict["name"] = datasource_name
+
+    id = "d53c2384-f973-4a0c-9c85-af1d67c06f58"
+    config_dict["id"] = id
+
+    datasource = context.add_or_update_datasource(**config_dict)
     # 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
+    assert datasource.config["id"] == 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
+    assert context.datasource_store.save_count == 1
 
 
 @pytest.mark.unit
+def test_add_or_update_datasource_updates_with_existing_datasource_successfully(
+    in_memory_data_context: EphemeralDataContextSpy,
+    datasource_name: str,
+    datasource_config: DatasourceConfig,
+):
+    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."
+
+    config_dict = datasource_config.to_dict()
+    config_dict["name"] = datasource_name
+    for attr in ("class_name", "module_name"):
+        config_dict.pop(attr)
+
+    datasource = Datasource(**config_dict)
+    persisted_datasource = context.add_or_update_datasource(datasource=datasource)
+
+    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.datasource_store.save_count == 1
+
+    expected_config = datasource.config
+    actual_config = persisted_datasource.config
+
+    # Name gets injected into data connector as part of serialization hook
+    # Removing for purposes of test assertion
+    data_connectors = actual_config["data_connectors"]
+    data_connector_name = tuple(data_connectors.keys())[0]
+    data_connectors[data_connector_name].pop("name")
+
+    assert actual_config == expected_config
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("use_existing_datasource", [True, False])
 def test_add_or_update_datasource_adds_successfully(
     in_memory_data_context: EphemeralDataContextSpy,
+    datasource_config: DatasourceConfig,
+    use_existing_datasource: bool,
 ):
     context = in_memory_data_context
 
@@ -289,11 +399,16 @@ def test_add_or_update_datasource_adds_successfully(
 
     assert datasource_name not in context.datasources
 
-    _ = context.add_or_update_datasource(
-        name=datasource_name,
-        module_name="great_expectations.datasource",
-        class_name="PandasDatasource",
-    )
+    config_dict = datasource_config.to_dict()
+    config_dict["name"] = datasource_name
+
+    if use_existing_datasource:
+        for attr in ("class_name", "module_name"):
+            config_dict.pop(attr)
+        datasource = Datasource(**config_dict)
+        _ = context.add_or_update_datasource(datasource=datasource)
+    else:
+        _ = context.add_or_update_datasource(**config_dict)
 
     num_datasource_after = len(context.datasources)
     num_datasource_configs_after = len(context.config.datasources)
@@ -301,7 +416,36 @@ def test_add_or_update_datasource_adds_successfully(
     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
+    assert context.datasource_store.save_count == 1
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "datasource,datasource_name",
+    [
+        pytest.param(
+            mock.MagicMock(),
+            "my_datasource",
+            id="both datasource and name",
+        ),
+        pytest.param(None, None, id="neither datasource nor name"),
+    ],
+)
+def test_add_or_update_datasource_conflicting_args_failure(
+    in_memory_data_context: EphemeralDataContextSpy,
+    datasource: mock.MagicMock | None,
+    datasource_name: str | None,
+):
+    context = in_memory_data_context
+
+    with pytest.raises(ValueError) as e:
+        context.add_or_update_datasource(datasource=datasource, name=datasource_name)
+
+    assert (
+        "an existing datasource or individual constructor arguments (but not both)"
+        in str(e.value)
+    )
+    assert context.expectations_store.save_count == 0
 
 
 @pytest.mark.unit
diff --git a/tests/data_context/test_loading_and_saving_of_data_context_configs.py b/tests/data_context/test_loading_and_saving_of_data_context_configs.py
index c5b4913fd..99658dcbe 100644
--- a/tests/data_context/test_loading_and_saving_of_data_context_configs.py
+++ b/tests/data_context/test_loading_and_saving_of_data_context_configs.py
@@ -30,7 +30,7 @@ def test_add_datasource(empty_data_context):
     config_filename = context.root_directory + "/great_expectations.yml"
 
     # Config can't be instantiated
-    with pytest.raises(TypeError):
+    with pytest.raises(gx_exceptions.DatasourceInitializationError):
         context.add_datasource(
             "my_new_datasource",
             **{