Implement the Custom Endpoint Plugin (#727)

Author: aaron-congo

aws_advanced_python_wrapper/allowed_and_blocked_hosts.py

@@ -0,0 +1,29 @@
+#   Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+#   Licensed under the Apache License, Version 2.0 (the "License").
+#   You may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+
+from typing import Optional, Set
+
+
+class AllowedAndBlockedHosts:
+    def __init__(self, allowed_host_ids: Optional[Set[str]], blocked_host_ids: Optional[Set[str]]):
+        self._allowed_host_ids = None if not allowed_host_ids else allowed_host_ids
+        self._blocked_host_ids = None if not blocked_host_ids else blocked_host_ids
+
+    @property
+    def allowed_host_ids(self) -> Optional[Set[str]]:
+        return self._allowed_host_ids
+
+    @property
+    def blocked_host_ids(self) -> Optional[Set[str]]:
+        return self._blocked_host_ids

aws_advanced_python_wrapper/aurora_connection_tracker_plugin.py

@@ -201,15 +201,15 @@ def _connect(self, host_info: HostInfo, connect_func: Callable):
 
     def execute(self, target: object, method_name: str, execute_func: Callable, *args: Any, **kwargs: Any) -> Any:
         if self._current_writer is None or self._need_update_current_writer:
-            self._current_writer = self._get_writer(self._plugin_service.hosts)
+            self._current_writer = self._get_writer(self._plugin_service.all_hosts)
             self._need_update_current_writer = False
 
         try:
             return execute_func()
 
         except Exception as e:
             # Check that e is a FailoverError and that the writer has changed
-            if isinstance(e, FailoverError) and self._get_writer(self._plugin_service.hosts) != self._current_writer:
+            if isinstance(e, FailoverError) and self._get_writer(self._plugin_service.all_hosts) != self._current_writer:
                 self._tracker.invalidate_all_connections(host_info=self._current_writer)
                 self._tracker.log_opened_connections()
                 self._need_update_current_writer = True

aws_advanced_python_wrapper/aurora_initial_connection_strategy_plugin.py

@@ -198,7 +198,7 @@ def _delay(self, delay_ms: int):
         sleep(delay_ms / 1000)
 
     def _get_writer(self) -> Optional[HostInfo]:
-        for host in self._plugin_service.hosts:
+        for host in self._plugin_service.all_hosts:
             if host.role == HostRole.WRITER:
                 return host
 
@@ -225,10 +225,10 @@ def init_host_provider(self, props: Properties, host_list_provider_service: Host
         init_host_provider_func(props)
 
     def _has_no_readers(self) -> bool:
-        if len(self._plugin_service.hosts) == 0:
+        if len(self._plugin_service.all_hosts) == 0:
             return False
 
-        for host in self._plugin_service.hosts:
+        for host in self._plugin_service.all_hosts:
             if host.role == HostRole.READER:
                 return False
 

aws_advanced_python_wrapper/aws_secrets_manager_plugin.py

@@ -35,6 +35,7 @@
 from aws_advanced_python_wrapper.utils.messages import Messages
 from aws_advanced_python_wrapper.utils.properties import (Properties,
                                                           WrapperProperties)
+from aws_advanced_python_wrapper.utils.region_utils import RegionUtils
 from aws_advanced_python_wrapper.utils.telemetry.telemetry import \
     TelemetryTraceLevel
 
@@ -63,6 +64,7 @@ def __init__(self, plugin_service: PluginService, props: Properties, session: Op
                 Messages.get_formatted("AwsSecretsManagerPlugin.MissingRequiredConfigParameter",
                                        WrapperProperties.SECRETS_MANAGER_SECRET_ID.name))
 
+        self._region_utils = RegionUtils()
         region: str = self._get_rds_region(secret_id, props)
 
         secrets_endpoint = WrapperProperties.SECRETS_MANAGER_ENDPOINT.get(props)
@@ -194,23 +196,22 @@ def _apply_secret_to_properties(self, properties: Properties):
             WrapperProperties.PASSWORD.set(properties, self._secret.password)
 
     def _get_rds_region(self, secret_id: str, props: Properties) -> str:
-        region: Optional[str] = props.get(WrapperProperties.SECRETS_MANAGER_REGION.name)
-        if not region:
-            match = search(self._SECRETS_ARN_PATTERN, secret_id)
-            if match:
-                region = match.group("region")
-            else:
-                raise AwsWrapperError(
-                    Messages.get_formatted("AwsSecretsManagerPlugin.MissingRequiredConfigParameter",
-                                           WrapperProperties.SECRETS_MANAGER_REGION.name))
-
         session = self._session if self._session else boto3.Session()
-        if region not in session.get_available_regions("rds"):
-            exception_message = "AwsSdk.UnsupportedRegion"
-            logger.debug(exception_message, region)
-            raise AwsWrapperError(Messages.get_formatted(exception_message, region))
+        region = self._region_utils.get_region(props, WrapperProperties.SECRETS_MANAGER_REGION.name, session=session)
+
+        if region:
+            return region
 
-        return region
+        match = search(self._SECRETS_ARN_PATTERN, secret_id)
+        if match:
+            region = match.group("region")
+
+        if region:
+            return self._region_utils.verify_region(region)
+        else:
+            raise AwsWrapperError(
+                Messages.get_formatted("AwsSecretsManagerPlugin.MissingRequiredConfigParameter",
+                                       WrapperProperties.SECRETS_MANAGER_REGION.name))
 
 
 class AwsSecretsManagerPluginFactory(PluginFactory):