App Config Provider - Provider Refactor (#43196)

* Code Cleanup

* Move validation to shared file

* Updating Header Check

* Update test_azureappconfigurationproviderbase.py

* moved async tests to aio folder

* post merge updates

Author: mrm9084

sdk/appconfiguration/azure-appconfiguration-provider/assets.json

@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "python",
   "TagPrefix": "python/appconfiguration/azure-appconfiguration-provider",
-  "Tag": "python/appconfiguration/azure-appconfiguration-provider_b13d43b82a"
+  "Tag": "python/appconfiguration/azure-appconfiguration-provider_ccc89e9eaa"
 }

sdk/appconfiguration/azure-appconfiguration-provider/azure/appconfiguration/provider/_azureappconfigurationprovider.py

@@ -14,13 +14,13 @@
     overload,
     List,
     Tuple,
-    TYPE_CHECKING,
 )
 from azure.appconfiguration import (  # type:ignore # pylint:disable=no-name-in-module
     ConfigurationSetting,
     FeatureFlagConfigurationSetting,
     SecretReferenceConfigurationSetting,
 )
+from azure.core.credentials import TokenCredential
 from azure.core.exceptions import AzureError, HttpResponseError
 from ._models import AzureAppConfigurationKeyVaultOptions, SettingSelector
 from ._key_vault._secret_provider import SecretProvider
@@ -30,16 +30,13 @@
 )
 from ._azureappconfigurationproviderbase import (
     AzureAppConfigurationProviderBase,
-    update_correlation_context_header,
     delay_failure,
+    process_load_parameters,
     sdk_allowed_kwargs,
 )
 from ._client_manager import ConfigurationClientManager, _ConfigurationClientWrapper as ConfigurationClient
 from ._user_agent import USER_AGENT
 
-if TYPE_CHECKING:
-    from azure.core.credentials import TokenCredential
-
 JSON = Mapping[str, Any]
 logger = logging.getLogger(__name__)
 
@@ -176,51 +173,19 @@ def load(  # pylint: disable=docstring-keyword-should-match-keyword-only
 
 
 def load(*args, **kwargs) -> "AzureAppConfigurationProvider":
-    endpoint: Optional[str] = kwargs.pop("endpoint", None)
-    credential: Optional["TokenCredential"] = kwargs.pop("credential", None)
-    connection_string: Optional[str] = kwargs.pop("connection_string", None)
-    key_vault_options: Optional[AzureAppConfigurationKeyVaultOptions] = kwargs.pop("key_vault_options", None)
     start_time = datetime.datetime.now()
 
-    # Update endpoint and credential if specified positionally.
-    if len(args) > 2:
-        raise TypeError(
-            "Unexpected positional parameters. Please pass either endpoint and credential, or a connection string."
-        )
-    if len(args) == 1:
-        if endpoint is not None:
-            raise TypeError("Received multiple values for parameter 'endpoint'.")
-        endpoint = args[0]
-    elif len(args) == 2:
-        if credential is not None:
-            raise TypeError("Received multiple values for parameter 'credential'.")
-        endpoint, credential = args
-
-    if (endpoint or credential) and connection_string:
-        raise ValueError("Please pass either endpoint and credential, or a connection string.")
-
-    # Removing use of AzureAppConfigurationKeyVaultOptions
-    if key_vault_options:
-        if "keyvault_credential" in kwargs or "secret_resolver" in kwargs or "keyvault_client_configs" in kwargs:
-            raise ValueError(
-                "Key Vault configurations should only be set by either the key_vault_options or kwargs not both."
-            )
-        kwargs["keyvault_credential"] = key_vault_options.credential
-        kwargs["secret_resolver"] = key_vault_options.secret_resolver
-        kwargs["keyvault_client_configs"] = key_vault_options.client_configs
-
-    if kwargs.get("keyvault_credential") is not None and kwargs.get("secret_resolver") is not None:
-        raise ValueError("A keyvault credential and secret resolver can't both be configured.")
-
-    uses_key_vault = (
-        "keyvault_credential" in kwargs
-        or "keyvault_client_configs" in kwargs
-        or "secret_resolver" in kwargs
-        or kwargs.get("uses_key_vault", False)
-    )
+    # Process common load parameters using shared logic
+    params = process_load_parameters(*args, **kwargs)
 
-    provider = _buildprovider(connection_string, endpoint, credential, uses_key_vault=uses_key_vault, **kwargs)
-    kwargs = sdk_allowed_kwargs(kwargs)
+    provider = _buildprovider(
+        params["connection_string"],
+        params["endpoint"],
+        params["credential"],
+        uses_key_vault=params["uses_key_vault"],
+        **params["kwargs"],
+    )
+    kwargs = sdk_allowed_kwargs(params["kwargs"])
 
     try:
         provider._load_all(**kwargs)  # pylint:disable=protected-access
@@ -231,7 +196,7 @@ def load(*args, **kwargs) -> "AzureAppConfigurationProvider":
 
 
 def _buildprovider(
-    connection_string: Optional[str], endpoint: Optional[str], credential: Optional["TokenCredential"], **kwargs
+    connection_string: Optional[str], endpoint: Optional[str], credential: Optional[TokenCredential], **kwargs
 ) -> "AzureAppConfigurationProvider":
     # pylint:disable=protected-access
     if connection_string:
@@ -254,7 +219,7 @@ class AzureAppConfigurationProvider(AzureAppConfigurationProviderBase):  # pylin
     """
 
     def __init__(self, **kwargs: Any) -> None:
-        super().__init__(**kwargs)
+        super(AzureAppConfigurationProvider, self).__init__(**kwargs)
 
         if "user_agent" in kwargs:
             user_agent = kwargs.pop("user_agent") + " " + USER_AGENT
@@ -288,17 +253,12 @@ def __init__(self, **kwargs: Any) -> None:
 
     def _attempt_refresh(self, client: ConfigurationClient, replica_count: int, is_failover_request: bool, **kwargs):
         settings_refreshed = False
-        headers = update_correlation_context_header(
+        headers = self._update_correlation_context_header(
             kwargs.pop("headers", {}),
             "Watch",
             replica_count,
-            self._feature_flag_enabled,
-            self._feature_filter_usage,
             self._secret_provider.uses_key_vault,
-            self._uses_load_balancing,
             is_failover_request,
-            self._uses_ai_configuration,
-            self._uses_aicc_configuration,
         )
         configuration_settings: List[ConfigurationSetting] = []
         feature_flags: Optional[List[FeatureFlagConfigurationSetting]] = None
@@ -341,17 +301,7 @@ def _attempt_refresh(self, client: ConfigurationClient, replica_count: int, is_f
                 # Configuration Settings have been refreshed
                 processed_settings = self._process_configurations(configuration_settings)
 
-            if feature_flags:
-                # Reset feature flag usage
-                self._feature_filter_usage = {}
-                processed_feature_flags = [self._process_feature_flag(ff) for ff in feature_flags]
-
-            if self._feature_flag_enabled:
-                # Create the feature management schema and add feature flags
-                if feature_flags:
-                    self._watched_feature_flags = self._update_watched_feature_flags(feature_flags)
-                processed_settings[FEATURE_MANAGEMENT_KEY] = {}
-                processed_settings[FEATURE_MANAGEMENT_KEY][FEATURE_FLAG_KEY] = processed_feature_flags
+            processed_settings = self._process_feature_flags(processed_settings, processed_feature_flags, feature_flags)
             self._dict = processed_settings
             if settings_refreshed:
                 # Update the watch keys that have changed
@@ -363,7 +313,6 @@ def _attempt_refresh(self, client: ConfigurationClient, replica_count: int, is_f
                 self._feature_flag_refresh_timer.reset()
             if (settings_refreshed or feature_flags) and self._on_refresh_success:
                 self._on_refresh_success()
-            return
         except AzureError as e:
             logger.warning("Failed to refresh configurations from endpoint %s", client.endpoint)
             self._replica_client_manager.backoff(client)
@@ -412,7 +361,7 @@ def refresh(self, **kwargs) -> None:
         finally:
             self._refresh_lock.release()
 
-    def _load_all(self, **kwargs):
+    def _load_all(self, **kwargs: Any) -> None:
         self._replica_client_manager.refresh_clients()
         self._replica_client_manager.find_active_clients()
         is_failover_request = False
@@ -424,21 +373,15 @@ def _load_all(self, **kwargs):
         exception: Exception = RuntimeError(error_message)
 
         while client := self._replica_client_manager.get_next_active_client():
-            headers = update_correlation_context_header(
+            headers = self._update_correlation_context_header(
                 kwargs.pop("headers", {}),
                 "Startup",
                 replica_count,
-                self._feature_flag_enabled,
-                self._feature_filter_usage,
                 self._secret_provider.uses_key_vault,
-                self._uses_load_balancing,
                 is_failover_request,
-                self._uses_ai_configuration,
-                self._uses_aicc_configuration,
             )
             try:
                 configuration_settings = client.load_configuration_settings(self._selects, headers=headers, **kwargs)
-                processed_feature_flags = []
                 watched_settings = self._update_watched_settings(configuration_settings)
                 processed_settings = self._process_configurations(configuration_settings)
 
@@ -448,12 +391,7 @@ def _load_all(self, **kwargs):
                         headers=headers,
                         **kwargs,
                     )
-                    processed_feature_flags = (
-                        [self._process_feature_flag(ff) for ff in feature_flags] if feature_flags else []
-                    )
-                    processed_settings[FEATURE_MANAGEMENT_KEY] = {}
-                    processed_settings[FEATURE_MANAGEMENT_KEY][FEATURE_FLAG_KEY] = processed_feature_flags
-                    self._watched_feature_flags = self._update_watched_feature_flags(feature_flags)
+                    processed_settings = self._process_feature_flags(processed_settings, [], feature_flags)
                 for (key, label), etag in self._watched_settings.items():
                     if not etag:
                         try: