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

aws_advanced_python_wrapper/stale_dns_plugin.py

@@ -0,0 +1,209 @@
+#  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
+
+import socket
+from typing import TYPE_CHECKING, Any, Callable, Dict, 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.hostinfo import HostInfo
+    from aws_advanced_python_wrapper.pep249 import Connection
+    from aws_advanced_python_wrapper.plugin_service import PluginService
+    from aws_advanced_python_wrapper.utils.properties import Properties
+
+from aws_advanced_python_wrapper.hostinfo import 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.messages import Messages
+from aws_advanced_python_wrapper.utils.notifications import HostEvent
+from aws_advanced_python_wrapper.utils.rdsutils import RdsUtils
+from aws_advanced_python_wrapper.utils.utils import LogUtils
+
+logger = Logger(__name__)
+
+
+class StaleDnsHelper:
+    RETRIES: int = 3
+
+    def __init__(self, plugin_service: PluginService) -> None:
+        self._plugin_service = plugin_service
+        self._rds_helper = RdsUtils()
+        self._writer_host_info: Optional[HostInfo] = None
+        self._writer_host_address: Optional[str] = None
+
+    def get_verified_connection(self, is_initial_connection: bool, host_list_provider_service: HostListProviderService, host_info: HostInfo,
+                                props: Properties, connect_func: Callable) -> Connection:
+        """
+        Ensure the connection created is not a stale writer connection that
+        :param is_initial_connection:
+        :param host_list_provider_service:
+        :param host_info:
+        :param props:
+        :param connect_func:
+        :return:
+        """
+        if not self._rds_helper.is_writer_cluster_dns(host_info.host):
+            return connect_func()
+
+        conn: Connection = connect_func()
+
+        cluster_inet_address: Optional[str] = None
+        try:
+            cluster_inet_address = socket.gethostbyname(host_info.host)
+        except socket.gaierror:
+            pass
+
+        host_inet_address: Optional[str] = cluster_inet_address
+
+        logger.debug("StaleDnsHelper.ClusterEndpointDns", host_info.host, host_inet_address)
+
+        if cluster_inet_address is None:
+            return conn
+
+        if self._plugin_service.get_host_role(conn) == HostRole.READER:
+            # This if-statement is only reached if the connection url is a writer cluster endpoint.
+            # If the new connection resolves to a reader instance, this means the topology is outdated.
+            # Force refresh to update the topology.
+            self._plugin_service.force_refresh_host_list(conn)
+        else:
+            self._plugin_service.refresh_host_list(conn)
+
+        logger.debug("LogUtils.Topology", LogUtils.log_topology(self._plugin_service.hosts))
+
+        if self._writer_host_info is None:
+            writer_candidate: Optional[HostInfo] = self._get_writer()
+            if writer_candidate is not None and self._rds_helper.is_rds_cluster_dns(writer_candidate.host):
+                return conn
+
+            self._writer_host_info = writer_candidate
+
+        logger.debug("StaleDnsHelper.WriterHostSpec", self._writer_host_info)
+
+        if self._writer_host_info is None:
+            return conn
+
+        if self._writer_host_address is None:
+            try:
+                self._writer_host_address = socket.gethostbyname(self._writer_host_info.host)
+            except socket.gaierror:
+                pass
+
+        logger.debug("StaleDnsHelper.WriterInetAddress", self._writer_host_address)
+
+        if self._writer_host_address is None:
+            return conn
+
+        if self._writer_host_address != cluster_inet_address:
+            logger.debug("StaleDnsHelper.StaleDnsDetected", self._writer_host_info)
+
+            writer_conn: Connection = self._plugin_service.connect(self._writer_host_info, props)
+            if is_initial_connection:
+                host_list_provider_service.initial_connection_host_info = self._writer_host_info
+
+            if conn is not None:
+                try:
+                    conn.close()
+                except Exception:
+                    pass
+                return writer_conn
+
+        return conn
+
+    def notify_host_list_changed(self, changes: Dict[str, Set[HostEvent]]) -> None:
+        if self._writer_host_info is None:
+            return
+
+        writer_changes = changes.get(self._writer_host_info.url, None)
+        if writer_changes is not None and HostEvent.CONVERTED_TO_READER in writer_changes:
+            logger.debug("StaleDnsHelper.Reset")
+            self._writer_host_info = None
+            self._writer_host_address = None
+
+    def _get_writer(self) -> Optional[HostInfo]:
+        for host in self._plugin_service.hosts:
+            if host.role == HostRole.WRITER:
+                return host
+        return None
+
+
+class StaleDnsPlugin(Plugin):
+
+    _SUBSCRIBED_METHODS: Set[str] = {"init_host_provider",
+                                     "connect",
+                                     "force_connect",
+                                     "notify_host_list_changed"}
+
+    def __init__(self, plugin_service: PluginService) -> None:
+        self._plugin_service = plugin_service
+        self._stale_dns_helper = StaleDnsHelper(self._plugin_service)
+
+        StaleDnsPlugin._SUBSCRIBED_METHODS.update(self._plugin_service.network_bound_methods)
+
+    @property
+    def subscribed_methods(self) -> Set[str]:
+        return self._SUBSCRIBED_METHODS
+
+    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._stale_dns_helper.get_verified_connection(
+            is_initial_connection, self._host_list_provider_service, host_info, props, 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._stale_dns_helper.get_verified_connection(
+            is_initial_connection, self._host_list_provider_service, host_info, props, force_connect_func)
+
+    def execute(self, target: type, method_name: str, execute_func: Callable, *args: Any, **kwargs: Any) -> Any:
+        try:
+            self._plugin_service.refresh_host_list()
+        except Exception:
+            pass
+
+        return execute_func()
+
+    def init_host_provider(
+            self,
+            properties: Properties,
+            host_list_provider_service: HostListProviderService,
+            init_host_provider_func: Callable):
+
+        self._host_list_provider_service = host_list_provider_service
+
+        init_host_provider_func()
+
+        if self._host_list_provider_service.is_static_host_list_provider():
+            raise Exception(Messages.get_formatted("StaleDnsPlugin.RequireDynamicProvider"))
+
+    def notify_host_list_changed(self, changes: Dict[str, Set[HostEvent]]):
+        self._stale_dns_helper.notify_host_list_changed(changes)
+
+
+class StaleDnsPluginFactory(PluginFactory):
+    def get_instance(self, plugin_service: PluginService, props: Properties) -> Plugin:
+        return StaleDnsPlugin(plugin_service)

aws_advanced_python_wrapper/states/__init__.py

@@ -0,0 +1,13 @@
+#  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.

aws_advanced_python_wrapper/states/session_state.py

@@ -0,0 +1,94 @@
+#  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 typing import Generic, Optional, TypeVar
+
+T = TypeVar("T")
+
+
+class SessionStateField(Generic[T]):
+    _value: Optional[T] = None
+    _pristine_value: Optional[T] = None
+
+    def copy(self):
+        new_field: SessionStateField[T] = SessionStateField()
+        if self._value is not None:
+            new_field._value = self._value
+        if self._pristine_value is not None:
+            new_field._pristine_value = self._pristine_value
+        return new_field
+
+    @property
+    def value(self):
+        return self._value
+
+    @value.setter
+    def value(self, new_value: Optional[T]):
+        self._value = new_value
+
+    @property
+    def pristine_value(self):
+        return self._pristine_value
+
+    @pristine_value.setter
+    def pristine_value(self, new_value: Optional[T]):
+        self._pristine_value = new_value
+
+    def reset_value(self):
+        self._value = None
+
+    def reset_pristine_value(self):
+        self.pristine_value = None
+
+    def reset(self):
+        self.reset_value()
+        self.reset_pristine_value()
+
+    def is_pristine(self) -> bool:
+        if self.value is None:
+            # the value has never been set up so the session state has pristine value
+            return True
+        return self._pristine_value is not None and self._pristine_value == self.value
+
+    def can_restore_pristine(self):
+        if self._pristine_value is None:
+            return False
+
+        if self._value is not None:
+            # it's necessary to restore pristine value only if current session value is not the same as pristine value.
+            return self.value != self.pristine_value
+
+        # it's inconclusive if the current value is the same as pristine value, so we need to take the safest path.
+        return True
+
+    def __str__(self):
+        return f"{self.pristine_value if self.pristine_value is not None else '(blank)'} -> {self.value if self.value is not None else '(blank)'}"
+
+    def __repr__(self):
+        return f"{self.pristine_value if self.pristine_value is not None else '(blank)'} -> {self.value if self.value is not None else '(blank)'}"
+
+
+class SessionState:
+    def __init__(self):
+        self.auto_commit: SessionStateField[bool] = SessionStateField()
+        self.readonly: SessionStateField[bool] = SessionStateField()
+
+    def copy(self):
+        new_session_state: SessionState = SessionState()
+        new_session_state.auto_commit = self.auto_commit.copy()
+        new_session_state.readonly = self.readonly.copy()
+
+        return new_session_state

aws_advanced_python_wrapper/states/session_state_service.py

@@ -0,0 +1,221 @@
+#  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 typing import TYPE_CHECKING, Callable, Optional, Protocol
+
+if TYPE_CHECKING:
+    from aws_advanced_python_wrapper.plugin_service import PluginService
+    from aws_advanced_python_wrapper.pep249 import Connection
+
+from aws_advanced_python_wrapper.errors import AwsWrapperError
+from aws_advanced_python_wrapper.states.session_state import SessionState
+from aws_advanced_python_wrapper.utils.log import Logger
+from aws_advanced_python_wrapper.utils.properties import (Properties,
+                                                          WrapperProperties)
+
+logger = Logger(__name__)
+
+
+class SessionStateService(Protocol):
+
+    def get_autocommit(self) -> Optional[bool]:
+        ...
+
+    def set_autocommit(self, autocommit: bool):
+        ...
+
+    def setup_pristine_autocommit(self, autocommit: Optional[bool] = None):
+        ...
+
+    def get_readonly(self):
+        ...
+
+    def set_read_only(self, readonly: bool):
+        ...
+
+    def setup_pristine_readonly(self, readonly: Optional[bool] = None):
+        ...
+
+    def reset(self):
+        ...
+
+    def begin(self):
+        """
+        Begin session transfer process
+        """
+        ...
+
+    def complete(self):
+        """
+        Complete session transfer process. This method should be called despite whether session transfer is successful or not.
+        """
+        ...
+
+    def apply_current_session_state(self, new_connection: Connection):
+        ...
+
+    def apply_pristine_session_state(self, new_connection: Connection):
+        ...
+
+
+class SessionStateTransferHandlers:
+    reset_session_state_on_close_callable: Optional[Callable] = None
+    transfer_session_state_on_switch_callable: Optional[Callable] = None
+
+    @staticmethod
+    def set_reset_session_state_on_close_func(func: Callable):
+        SessionStateTransferHandlers.reset_session_state_on_close_callable = func
+
+    @staticmethod
+    def clear_reset_session_state_on_close_func():
+        SessionStateTransferHandlers.reset_session_state_on_close_callable = None
+
+    @staticmethod
+    def get_reset_session_state_on_close_func() -> Optional[Callable]:
+        return SessionStateTransferHandlers.reset_session_state_on_close_callable
+
+    @staticmethod
+    def get_transfer_session_state_on_switch_func() -> Optional[Callable]:
+        return SessionStateTransferHandlers.transfer_session_state_on_switch_callable
+
+    @staticmethod
+    def set_transfer_session_state_on_switch_func(func: Callable):
+        SessionStateTransferHandlers.transfer_session_state_on_switch_callable = func
+
+    @staticmethod
+    def reset_transfer_session_state_on_switch_func():
+        SessionStateTransferHandlers.reset_transfer_session_state_on_switch_callable = None
+
+
+class SessionStateServiceImpl(SessionStateService):
+    def __init__(self, plugin_service: PluginService, props: Properties):
+        self._session_state: SessionState = SessionState()
+        self._copy_session_state: Optional[SessionState] = None
+        self._plugin_service: PluginService = plugin_service
+        self._props: Properties = props
+
+    def log_current_state(self):
+        logger.debug(f"Current session state: \n{self._session_state}")
+
+    def _transfer_state_enabled_setting(self) -> bool:
+        return WrapperProperties.TRANSFER_SESSION_STATE_ON_SWITCH.get_bool(self._props)
+
+    def _reset_state_enabled_setting(self) -> bool:
+        return WrapperProperties.RESET_SESSION_STATE_ON_CLOSE.get_bool(self._props)
+
+    def get_autocommit(self) -> Optional[bool]:
+        return self._session_state.auto_commit.value
+
+    def set_autocommit(self, autocommit: bool):
+        if not self._transfer_state_enabled_setting():
+            return
+        self._session_state.auto_commit.value = autocommit
+
+    def setup_pristine_autocommit(self, autocommit: Optional[bool] = None):
+        if not self._transfer_state_enabled_setting():
+            return
+        if self._session_state.auto_commit.pristine_value is not None:
+            return
+
+        if autocommit is None and self._plugin_service.current_connection is not None:
+            autocommit = self._plugin_service.driver_dialect.get_autocommit(self._plugin_service.current_connection)
+
+        self._session_state.auto_commit.pristine_value = autocommit
+        self.log_current_state()
+
+    def get_readonly(self):
+        return self._session_state.readonly.value
+
+    def set_read_only(self, readonly: bool):
+        if not self._transfer_state_enabled_setting():
+            return
+        self._session_state.readonly.value = readonly
+
+    def setup_pristine_readonly(self, readonly: Optional[bool] = None):
+        if not self._transfer_state_enabled_setting():
+            return
+        if self._session_state.readonly.pristine_value is not None:
+            return
+
+        if readonly is None and self._plugin_service.current_connection is not None:
+            readonly = self._plugin_service.driver_dialect.is_read_only(self._plugin_service.current_connection)
+
+        self._session_state.readonly.pristine_value = readonly
+        self.log_current_state()
+
+    def reset(self):
+        self._session_state.auto_commit.reset()
+        self._session_state.readonly.reset()
+
+    def begin(self):
+        self.log_current_state()
+        if not self._transfer_state_enabled_setting() and not self._reset_state_enabled_setting():
+            return
+        if self._copy_session_state is not None:
+            raise AwsWrapperError("Previous session state transfer is not completed.")
+        self._copy_session_state = self._session_state.copy()
+
+    def complete(self):
+        self._copy_session_state = None
+
+    def apply_current_session_state(self, new_connection: Connection):
+        if not self._transfer_state_enabled_setting():
+            return
+
+        func: Optional[Callable] = SessionStateTransferHandlers.get_transfer_session_state_on_switch_func()
+        if func is not None:
+            is_handled: bool = func(self._session_state, new_connection)
+            if is_handled:
+                # Custom function has handled session transfer
+                return
+
+        if self._session_state.auto_commit.value is not None:
+            self._session_state.auto_commit.reset_pristine_value()
+            self.setup_pristine_autocommit()
+            self._plugin_service.driver_dialect.set_autocommit(new_connection, self._session_state.auto_commit.value)
+
+        if self._session_state.readonly.value is not None:
+            self._session_state.readonly.reset_pristine_value()
+            self.setup_pristine_readonly()
+            self._plugin_service.driver_dialect.set_read_only(new_connection, self._session_state.readonly.value)
+
+    def apply_pristine_session_state(self, new_connection: Connection):
+        if not self._transfer_state_enabled_setting():
+            return
+
+        func: Optional[Callable] = SessionStateTransferHandlers.get_transfer_session_state_on_switch_func()
+        if func is not None:
+            is_handled: bool = func(self._session_state, new_connection)
+            if is_handled:
+                # Custom function has handled session transfer
+                return
+        if self._copy_session_state is None:
+            return
+
+        if self._copy_session_state.auto_commit.can_restore_pristine():
+            try:
+                self._plugin_service.driver_dialect.set_autocommit(new_connection, self._copy_session_state.auto_commit.pristine_value)
+            except Exception:
+                # Ignore any exception.
+                pass
+
+        if self._copy_session_state.readonly.can_restore_pristine():
+            try:
+                self._plugin_service.driver_dialect.set_read_only(new_connection, self._copy_session_state.readonly.pristine_value)
+            except Exception:
+                # Ignore any exception.
+                pass

aws_advanced_python_wrapper/utils/__init__.py

@@ -0,0 +1,13 @@
+#  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.

aws_advanced_python_wrapper/utils/atomic.py

@@ -0,0 +1,51 @@
+#  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 threading import Lock
+
+
+class AtomicInt:
+    def __init__(self, initial_value: int = 0):
+        self._value = initial_value
+        self._lock: Lock = Lock()
+
+    def get(self):
+        with self._lock:
+            return self._value
+
+    def set(self, value: int):
+        with self._lock:
+            self._value = value
+
+    def get_and_increment(self):
+        with self._lock:
+            value = self._value
+            self._value += 1
+            return value
+
+    def increment_and_get(self):
+        with self._lock:
+            self._value += 1
+            return self._value
+
+    def get_and_decrement(self):
+        with self._lock:
+            value = self._value
+            self._value -= 1
+            return value
+
+    def decrement_and_get(self):
+        with self._lock:
+            self._value -= 1
+            return self._value