aws-advanced-python-wrapper 1.0.0__py3-none-any.whl

CONTRIBUTING.md

@@ -0,0 +1,63 @@
+# Contributing Guidelines
+
+Thank you for your interest in contributing to our project. Whether it's a bug report, new feature, correction, or additional
+documentation, we greatly value feedback and contributions from our community.
+
+Please read through this document before submitting any issues or pull requests to ensure we have all the necessary
+information to effectively respond to your bug report or contribution.
+
+
+## Reporting Bugs/Feature Requests
+
+We welcome you to use the GitHub issue tracker to report bugs or suggest features.
+
+When filing an issue, please check existing open, or recently closed, issues to make sure somebody else hasn't already
+reported the issue. Please try to include as much information as you can. Details like these are incredibly useful:
+
+* A reproducible test case or series of steps
+* The version of our code being used
+* Any modifications you've made relevant to the bug
+* Anything unusual about your environment or deployment
+
+
+## Contributing via Pull Requests
+
+Contributions via pull requests are much appreciated. Before sending us a pull request, please ensure that:
+
+1. You are working against the latest source on the *main* branch.
+2. You check existing open, and recently merged, pull requests to make sure someone else hasn't addressed the problem already.
+3. You open an issue to discuss any significant work - we would hate for your time to be wasted.
+
+To send us a pull request, please:
+
+1. Fork the repository.
+2. Modify the source; please focus on the specific change you are contributing. If you also reformat all the code, it will be hard for us to focus on your change.
+3. Ensure local tests pass.
+4. Commit to your fork using clear commit messages.
+5. Send us a pull request, answering any default questions in the pull request interface.
+6. Pay attention to any automated CI failures reported in the pull request, and stay involved in the conversation.
+
+GitHub provides additional document on [forking a repository](https://help.github.com/articles/fork-a-repo/) and
+[creating a pull request](https://help.github.com/articles/creating-a-pull-request/).
+
+
+## Finding contributions to work on
+
+Looking at the existing issues is a great way to find something to contribute on. As our projects, by default, use the default GitHub issue labels (enhancement/bug/duplicate/help wanted/invalid/question/wontfix), looking at any 'help wanted' issues is a great place to start.
+
+
+## Code of Conduct
+
+This project has adopted the [Amazon Open Source Code of Conduct](https://aws.github.io/code-of-conduct).
+For more information see the [Code of Conduct FAQ](https://aws.github.io/code-of-conduct-faq) or contact
+opensource-codeofconduct@amazon.com with any additional questions or comments.
+
+
+## Security issue notifications
+
+If you discover a potential security issue in this project we ask that you notify AWS/Amazon Security via our [vulnerability reporting page](http://aws.amazon.com/security/vulnerability-reporting/). Please do **not** create a public GitHub issue.
+
+
+## Licensing
+
+See the [LICENSE](LICENSE) file for our project's licensing. We will ask you to confirm the licensing of your contribution.

aws_advanced_python_wrapper/__init__.py

@@ -0,0 +1,28 @@
+#  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 logging import DEBUG, getLogger
+
+from .utils.utils import LogUtils
+from .wrapper import AwsWrapperConnection
+
+# PEP249 compliance
+connect = AwsWrapperConnection.connect
+apilevel = "2.0"
+threadsafety = 2
+paramstyle = "pyformat"
+
+
+def set_logger(name='aws_advanced_python_wrapper', level=DEBUG, format_string=None):
+    LogUtils.setup_logger(getLogger(name), level, format_string)

aws_advanced_python_wrapper/aurora_connection_tracker_plugin.py

@@ -0,0 +1,228 @@
+#  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 __future__ import annotations
+
+from threading import Thread
+from typing import (TYPE_CHECKING, Any, Callable, Dict, FrozenSet, Optional,
+                    Set, Tuple)
+
+if TYPE_CHECKING:
+    from aws_advanced_python_wrapper.driver_dialect import DriverDialect
+    from aws_advanced_python_wrapper.plugin_service import PluginService
+    from aws_advanced_python_wrapper.pep249 import Connection
+
+    from aws_advanced_python_wrapper.utils.rds_url_type import RdsUrlType
+    from aws_advanced_python_wrapper.utils.properties import Properties
+
+from _weakrefset import WeakSet
+
+from aws_advanced_python_wrapper.errors import FailoverError
+from aws_advanced_python_wrapper.hostinfo import HostInfo, HostRole
+from aws_advanced_python_wrapper.plugin import Plugin, PluginFactory
+from aws_advanced_python_wrapper.utils.log import Logger
+from aws_advanced_python_wrapper.utils.rdsutils import RdsUtils
+
+logger = Logger(__name__)
+
+
+class OpenedConnectionTracker:
+    _opened_connections: Dict[str, WeakSet] = {}
+    _rds_utils = RdsUtils()
+
+    def populate_opened_connection_set(self, host_info: HostInfo, conn: Connection):
+        """
+        Add the given connection to the set of tracked connections.
+
+        :param host_info: host information of the given connection.
+        :param conn: currently opened connection.
+        """
+
+        aliases: FrozenSet[str] = host_info.as_aliases()
+        host: str = host_info.as_alias()
+
+        if self._rds_utils.is_rds_instance(host):
+            self._track_connection(host, conn)
+            return
+
+        instance_endpoint: Optional[str] = next((alias for alias in aliases if self._rds_utils.is_rds_instance(alias)),
+                                                None)
+        if not instance_endpoint:
+            logger.debug("OpenedConnectionTracker.UnableToPopulateOpenedConnectionSet")
+            return
+
+        self._track_connection(instance_endpoint, conn)
+
+    def invalidate_all_connections(self, host_info: Optional[HostInfo] = None, host: Optional[FrozenSet[str]] = None):
+        """
+        Invalidates all opened connections pointing to the same host in a daemon thread.
+
+       :param host_info: the :py:class:`HostInfo` object containing the URL of the host.
+       :param host: the set of aliases representing a specific host.
+        """
+
+        if host_info:
+            self.invalidate_all_connections(host=frozenset(host_info.as_alias()))
+            self.invalidate_all_connections(host=host_info.as_aliases())
+            return
+
+        instance_endpoint: Optional[str] = None
+        if host is None:
+            return
+
+        for instance in host:
+            if instance is not None and self._rds_utils.is_rds_instance(instance):
+                instance_endpoint = instance
+                break
+
+        if not instance_endpoint:
+            return
+
+        connection_set: Optional[WeakSet] = self._opened_connections.get(instance_endpoint)
+        if connection_set is not None:
+            self._log_connection_set(instance_endpoint, connection_set)
+            self._invalidate_connections(connection_set)
+
+    def _track_connection(self, instance_endpoint: str, conn: Connection):
+        connection_set: Optional[WeakSet] = self._opened_connections.get(instance_endpoint)
+        if connection_set is None:
+            connection_set = WeakSet()
+            connection_set.add(conn)
+            self._opened_connections[instance_endpoint] = connection_set
+        else:
+            connection_set.add(conn)
+
+        self.log_opened_connections()
+
+    @staticmethod
+    def _task(connection_set: WeakSet):
+        while connection_set is not None and len(connection_set) > 0:
+            conn_reference = connection_set.pop()
+
+            if conn_reference is None:
+                continue
+
+            try:
+                conn_reference.close()
+            except Exception:
+                # Swallow this exception, current connection should be useless anyway
+                pass
+
+    def _invalidate_connections(self, connection_set: WeakSet):
+        invalidate_connection_thread: Thread = Thread(daemon=True, target=self._task,
+                                                      args=[connection_set])  # type: ignore
+        invalidate_connection_thread.start()
+
+    def log_opened_connections(self):
+        msg = ""
+        for key, conn_set in self._opened_connections.items():
+            conn = ""
+            for item in list(conn_set):
+                conn += f"\n\t\t{item}"
+
+            msg += f"\t[{key} : {conn}]"
+
+        return logger.debug("OpenedConnectionTracker.OpenedConnectionsTracked", msg)
+
+    def _log_connection_set(self, host: str, conn_set: Optional[WeakSet]):
+        if conn_set is None or len(conn_set) == 0:
+            return
+
+        conn = ""
+        for item in list(conn_set):
+            conn += f"\n\t\t{item}"
+
+        msg = host + f"[{conn}\n]"
+        logger.debug("OpenedConnectionTracker.InvalidatingConnections", msg)
+
+
+class AuroraConnectionTrackerPlugin(Plugin):
+    _SUBSCRIBED_METHODS: Set[str] = {"*"}
+    _current_writer: Optional[HostInfo] = None
+    _need_update_current_writer: bool = False
+
+    @property
+    def subscribed_methods(self) -> Set[str]:
+        return self._SUBSCRIBED_METHODS
+
+    def __init__(self,
+                 plugin_service: PluginService,
+                 props: Properties,
+                 rds_utils: RdsUtils = RdsUtils(),
+                 tracker: OpenedConnectionTracker = OpenedConnectionTracker()):
+        self._plugin_service = plugin_service
+        self._props = props
+        self._rds_utils = rds_utils
+        self._tracker = tracker
+
+    def connect(
+            self,
+            target_driver_func: Callable,
+            driver_dialect: DriverDialect,
+            host_info: HostInfo,
+            props: Properties,
+            is_initial_connection: bool,
+            connect_func: Callable) -> Connection:
+        return self._connect(host_info, connect_func)
+
+    def force_connect(
+            self,
+            target_driver_func: Callable,
+            driver_dialect: DriverDialect,
+            host_info: HostInfo,
+            props: Properties,
+            is_initial_connection: bool,
+            force_connect_func: Callable) -> Connection:
+        return self._connect(host_info, force_connect_func)
+
+    def _connect(self, host_info: HostInfo, connect_func: Callable):
+        conn = connect_func()
+
+        if conn:
+            url_type: RdsUrlType = self._rds_utils.identify_rds_type(host_info.host)
+            if url_type.is_rds_cluster:
+                host_info.reset_aliases()
+                self._plugin_service.fill_aliases(conn, host_info)
+
+            self._tracker.populate_opened_connection_set(host_info, conn)
+            self._tracker.log_opened_connections()
+
+        return conn
+
+    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._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:
+                self._tracker.invalidate_all_connections(host_info=self._current_writer)
+                self._tracker.log_opened_connections()
+                self._need_update_current_writer = True
+            raise e
+
+    def _get_writer(self, hosts: Tuple[HostInfo, ...]) -> Optional[HostInfo]:
+        for host in hosts:
+            if host.role == HostRole.WRITER:
+                return host
+        return None
+
+
+class AuroraConnectionTrackerPluginFactory(PluginFactory):
+    def get_instance(self, plugin_service: PluginService, props: Properties) -> Plugin:
+        return AuroraConnectionTrackerPlugin(plugin_service, props)

aws_advanced_python_wrapper/aurora_initial_connection_strategy_plugin.py

@@ -0,0 +1,240 @@
+#  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 __future__ import annotations
+
+from time import perf_counter_ns, sleep
+from typing import TYPE_CHECKING, Callable, Optional, Set
+
+if TYPE_CHECKING:
+    from aws_advanced_python_wrapper.driver_dialect import DriverDialect
+    from aws_advanced_python_wrapper.host_list_provider import HostListProviderService
+    from aws_advanced_python_wrapper.pep249 import Connection
+    from aws_advanced_python_wrapper.plugin_service import PluginService
+
+from aws_advanced_python_wrapper.errors import AwsWrapperError
+from aws_advanced_python_wrapper.host_availability import HostAvailability
+from aws_advanced_python_wrapper.hostinfo import HostInfo, HostRole
+from aws_advanced_python_wrapper.plugin import Plugin, PluginFactory
+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.rds_url_type import RdsUrlType
+from aws_advanced_python_wrapper.utils.rdsutils import RdsUtils
+
+
+class AuroraInitialConnectionStrategyPlugin(Plugin):
+    _SUBSCRIBED_METHODS: Set[str] = {"init_host_provider",
+                                     "connect",
+                                     "force_connect"}
+
+    _host_list_provider_service: Optional[HostListProviderService] = None
+
+    @property
+    def subscribed_methods(self) -> Set[str]:
+        return AuroraInitialConnectionStrategyPlugin._SUBSCRIBED_METHODS
+
+    def __init__(self, plugin_service: PluginService):
+        super()
+        self._plugin_service = plugin_service
+        self._rds_utils = RdsUtils()
+
+    def connect(self, target_driver_func: Callable, driver_dialect: DriverDialect, host_info: HostInfo, props: Properties,
+                is_initial_connection: bool, connect_func: Callable) -> Connection:
+        return self._connect(host_info, props, is_initial_connection, connect_func)
+
+    def force_connect(self, target_driver_func: Callable, driver_dialect: DriverDialect, host_info: HostInfo, props: Properties,
+                      is_initial_connection: bool, force_connect_func: Callable) -> Connection:
+        return self._connect(host_info, props, is_initial_connection, force_connect_func)
+
+    def _connect(self, host_info: HostInfo, props: Properties, is_initial_connection: bool, connect_func: Callable):
+        type: RdsUrlType = self._rds_utils.identify_rds_type(host_info.host)
+        if not type.is_rds_cluster:
+            return connect_func()
+
+        if type == RdsUrlType.RDS_WRITER_CLUSTER:
+            writer_candidate_conn: Optional[Connection] = self._get_verified_writer_connection(props, is_initial_connection, connect_func)
+            if writer_candidate_conn is None:
+                return connect_func()
+            return writer_candidate_conn
+
+        if type == RdsUrlType.RDS_READER_CLUSTER:
+            reader_candidate_conn: Optional[Connection] = self._get_verified_reader_connection(props, is_initial_connection, connect_func)
+            if reader_candidate_conn is None:
+                return connect_func()
+            return reader_candidate_conn
+
+    def _get_verified_writer_connection(self, props: Properties, is_initial_connection: bool, connect_func: Callable) -> Connection | None:
+        retry_delay_ms: int = WrapperProperties.OPEN_CONNECTION_RETRY_INTERVAL_MS.get_int(props)
+        end_time_nano = perf_counter_ns() + (WrapperProperties.OPEN_CONNECTION_RETRY_INTERVAL_MS.get_int(props) * 1000000)
+
+        writer_candidate_conn: Optional[Connection]
+        writer_candidate: Optional[HostInfo]
+
+        while perf_counter_ns() < end_time_nano:
+            writer_candidate_conn = None
+            writer_candidate = None
+
+            try:
+                writer_candidate = self._get_writer()
+                if writer_candidate is None or self._rds_utils.is_rds_cluster_dns(writer_candidate.host):
+                    # Writer is not found. Topology is outdated.
+                    writer_candidate_conn = connect_func()
+                    self._plugin_service.force_refresh_host_list(writer_candidate_conn)
+                    writer_candidate = self._plugin_service.identify_connection(writer_candidate_conn)
+
+                    if writer_candidate is not None and writer_candidate.role != HostRole.WRITER:
+                        self._close_connection(writer_candidate_conn)
+                        self._delay(retry_delay_ms)
+                        continue
+
+                    if is_initial_connection and self._host_list_provider_service is not None:
+                        self._host_list_provider_service.initial_connection_host_info = writer_candidate
+
+                    return writer_candidate_conn
+
+                writer_candidate_conn = self._plugin_service.connect(writer_candidate, props)
+
+                if self._plugin_service.get_host_role(writer_candidate_conn) != HostRole.WRITER:
+                    self._plugin_service.force_refresh_host_list(writer_candidate_conn)
+                    self._close_connection(writer_candidate_conn)
+                    self._delay(retry_delay_ms)
+                    continue
+
+                # Writer connection is valid and verified.
+                if is_initial_connection and self._host_list_provider_service is not None:
+                    self._host_list_provider_service.initial_connection_host_info = writer_candidate
+                return writer_candidate_conn
+            except Exception as e:
+                self._close_connection(writer_candidate_conn)
+                raise e
+
+        return None
+
+    def _get_verified_reader_connection(self, props: Properties, is_initial_connection: bool, connect_func: Callable) -> Optional[Connection]:
+        retry_delay_ms: int = WrapperProperties.OPEN_CONNECTION_RETRY_INTERVAL_MS.get_int(props)
+        end_time_nano = perf_counter_ns() + (WrapperProperties.OPEN_CONNECTION_RETRY_INTERVAL_MS.get_int(props) * 1000000)
+
+        reader_candidate_conn: Optional[Connection]
+        reader_candidate: Optional[HostInfo]
+
+        while perf_counter_ns() < end_time_nano:
+            reader_candidate_conn = None
+            reader_candidate = None
+
+            try:
+                reader_candidate = self._get_reader(props)
+                if reader_candidate is None or self._rds_utils.is_rds_cluster_dns(reader_candidate.host):
+                    # READER is not found. Topology is outdated.
+                    reader_candidate_conn = connect_func()
+                    self._plugin_service.force_refresh_host_list(reader_candidate_conn)
+                    reader_candidate = self._plugin_service.identify_connection(reader_candidate_conn)
+
+                    if reader_candidate is not None and reader_candidate.role != HostRole.READER:
+                        if self._has_no_readers():
+                            # Cluster has no readers. Simulate Aurora reader cluster endpoint logic and return the current writer connection.
+                            if is_initial_connection and self._host_list_provider_service is not None:
+                                self._host_list_provider_service.initial_connection_host_info = reader_candidate
+                            return reader_candidate_conn
+
+                        self._close_connection(reader_candidate_conn)
+                        self._delay(retry_delay_ms)
+                        continue
+
+                    if is_initial_connection and self._host_list_provider_service is not None:
+                        self._host_list_provider_service.initial_connection_host_info = reader_candidate
+                    return reader_candidate_conn
+
+                reader_candidate_conn = self._plugin_service.connect(reader_candidate, props)
+
+                if self._plugin_service.get_host_role(reader_candidate_conn) != HostRole.READER:
+                    # If the new connection resolves to a writer instance, the topology is outdated.
+                    # Force refresh to update the topology.
+                    self._plugin_service.force_refresh_host_list(reader_candidate_conn)
+
+                    if self._has_no_readers():
+                        # Cluster has no readers. Simulate Aurora reader cluster endpoint logic and return the current writer connection.
+                        if is_initial_connection and self._host_list_provider_service is not None:
+                            self._host_list_provider_service.initial_connection_host_info = reader_candidate
+                        return reader_candidate_conn
+
+                    self._close_connection(reader_candidate_conn)
+                    self._delay(retry_delay_ms)
+                    continue
+
+                # Reader connection is valid and verified.
+                if is_initial_connection and self._host_list_provider_service is not None:
+                    self._host_list_provider_service.initial_connection_host_info = reader_candidate
+                return reader_candidate_conn
+            except Exception as e:
+                self._close_connection(reader_candidate_conn)
+                if not self._plugin_service.is_login_exception(e) and reader_candidate is not None:
+                    self._plugin_service.set_availability(reader_candidate.as_aliases(), HostAvailability.UNAVAILABLE)
+
+                raise e
+
+        return None
+
+    def _close_connection(self, connection: Optional[Connection]):
+        if connection is not None:
+            try:
+                connection.close()
+            except Exception:
+                # ignore
+                pass
+
+    def _delay(self, delay_ms: int):
+        sleep(delay_ms / 1000)
+
+    def _get_writer(self) -> Optional[HostInfo]:
+        for host in self._plugin_service.hosts:
+            if host.role == HostRole.WRITER:
+                return host
+
+        return None
+
+    def _get_reader(self, props: Properties) -> Optional[HostInfo]:
+        strategy = WrapperProperties.READER_INITIAL_HOST_SELECTOR_STRATEGY.get(props)
+        if (self._plugin_service is not None
+                and strategy is not None
+                and self._plugin_service.accepts_strategy(HostRole.READER, strategy)):
+            try:
+                return self._plugin_service.get_host_info_by_strategy(HostRole.READER, strategy)
+            except Exception:
+                # Host isn't found.
+                return None
+
+        raise AwsWrapperError(Messages.get_formatted("AuroraInitialConnectionStrategyPlugin.UnsupportedStrategy", strategy))
+
+    def init_host_provider(self, props: Properties, host_list_provider_service: HostListProviderService, init_host_provider_func: Callable):
+        self._host_list_provider_service = host_list_provider_service
+        if host_list_provider_service.is_static_host_list_provider():
+            raise AwsWrapperError(Messages.get("AuroraInitialConnectionStrategyPlugin.RequireDynamicProvider"))
+
+        init_host_provider_func(props)
+
+    def _has_no_readers(self) -> bool:
+        if len(self._plugin_service.hosts) == 0:
+            return False
+
+        for host in self._plugin_service.hosts:
+            if host.role == HostRole.READER:
+                return False
+
+        return True
+
+
+class AuroraInitialConnectionStrategyPluginFactory(PluginFactory):
+    def get_instance(self, plugin_service: PluginService, props: Properties) -> Plugin:
+        return AuroraInitialConnectionStrategyPlugin(plugin_service)