read_write_splitting_plugin.py

#  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 copy import deepcopy
from typing import TYPE_CHECKING, Any, Callable, Optional, Set, Tuple

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.utils.properties import Properties
    from aws_advanced_python_wrapper.connection_provider import ConnectionProviderManager

from aws_advanced_python_wrapper.errors import (AwsWrapperError, FailoverError,
                                                ReadWriteSplittingError)
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.messages import Messages
from aws_advanced_python_wrapper.utils.notifications import (
    ConnectionEvent, OldConnectionSuggestedAction)
from aws_advanced_python_wrapper.utils.properties import WrapperProperties

logger = Logger(__name__)


class ReadWriteSplittingPlugin(Plugin):
    _SUBSCRIBED_METHODS: Set[str] = {"init_host_provider",
                                     "connect",
                                     "notify_connection_changed",
                                     "Connection.set_read_only"}
    _POOL_PROVIDER_CLASS_NAME = "aws_advanced_python_wrapper.sql_alchemy_connection_provider.SqlAlchemyPooledConnectionProvider"

    def __init__(self, plugin_service: PluginService, props: Properties):
        self._plugin_service = plugin_service
        self._properties = props
        self._host_list_provider_service: HostListProviderService
        self._writer_connection: Optional[Connection] = None
        self._reader_connection: Optional[Connection] = None
        self._reader_host_info: Optional[HostInfo] = None
        self._conn_provider_manager: ConnectionProviderManager = self._plugin_service.get_connection_provider_manager()
        self._is_reader_conn_from_internal_pool: bool = False
        self._is_writer_conn_from_internal_pool: bool = False
        self._in_read_write_split: bool = False

        self._reader_selector_strategy: str = ""
        strategy = WrapperProperties.READER_HOST_SELECTOR_STRATEGY.get(self._properties)
        if strategy is not None:
            self._reader_selector_strategy = strategy
        else:
            default_strategy = WrapperProperties.READER_HOST_SELECTOR_STRATEGY.default_value
            if default_strategy is not None:
                self._reader_selector_strategy = default_strategy

    @property
    def subscribed_methods(self) -> Set[str]:
        return self._SUBSCRIBED_METHODS

    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
        init_host_provider_func()

    def connect(
            self,
            target_driver_func: Callable,
            driver_dialect: DriverDialect,
            host_info: HostInfo,
            props: Properties,
            is_initial_connection: bool,
            connect_func: Callable) -> Connection:
        if not self._plugin_service.accepts_strategy(host_info.role, self._reader_selector_strategy):
            raise AwsWrapperError(
                Messages.get_formatted("ReadWriteSplittingPlugin.UnsupportedHostInfoSelectorStrategy",
                                       self._reader_selector_strategy))

        return self.connect_internal(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_internal(is_initial_connection, force_connect_func)

    def connect_internal(self, is_initial_connection: bool, connect_func: Callable) -> Connection:
        current_conn = connect_func()

        if not is_initial_connection or self._host_list_provider_service.is_static_host_list_provider():
            return current_conn

        current_role = self._plugin_service.get_host_role(current_conn)
        if current_role is None or current_role == HostRole.UNKNOWN:
            self._log_and_raise_exception("ReadWriteSplittingPlugin.ErrorVerifyingInitialHostSpecRole")

        current_host = self._plugin_service.initial_connection_host_info
        if current_host is not None:
            if current_role == current_host.role:
                return current_conn

            updated_host = deepcopy(current_host)
            updated_host.role = current_role
            self._host_list_provider_service.initial_connection_host_info = updated_host

        return current_conn

    def notify_connection_changed(self, changes: Set[ConnectionEvent]) -> OldConnectionSuggestedAction:
        self._update_internal_connection_info()

        if self._in_read_write_split:
            return OldConnectionSuggestedAction.PRESERVE

        return OldConnectionSuggestedAction.NO_OPINION

    def execute(self, target: type, method_name: str, execute_func: Callable, *args: Any, **kwargs: Any) -> Any:
        driver_dialect = self._plugin_service.driver_dialect
        conn: Optional[Connection] = driver_dialect.get_connection_from_obj(target)
        current_conn: Optional[Connection] = driver_dialect.unwrap_connection(self._plugin_service.current_connection)

        if conn is not None and conn != current_conn:
            msg = Messages.get_formatted("PluginManager.MethodInvokedAgainstOldConnection", target)
            raise AwsWrapperError(msg)

        if method_name == "Connection.set_read_only" and args is not None and len(args) > 0:
            self._switch_connection_if_required(args[0])

        try:
            return execute_func()
        except Exception as ex:
            if isinstance(ex, FailoverError):
                logger.debug("ReadWriteSplittingPlugin.FailoverExceptionWhileExecutingCommand", method_name)
                self._close_idle_connections()
            else:
                logger.debug("ReadWriteSplittingPlugin.ExceptionWhileExecutingCommand", method_name)
            raise ex

    def _update_internal_connection_info(self):
        current_conn = self._plugin_service.current_connection
        current_host = self._plugin_service.current_host_info
        if current_conn is None or current_host is None:
            return

        if current_host.role == HostRole.WRITER:
            self._set_writer_connection(current_conn, current_host)
        else:
            self._set_reader_connection(current_conn, current_host)

    def _set_writer_connection(self, writer_conn: Connection, writer_host_info: HostInfo):
        self._writer_connection = writer_conn
        logger.debug("ReadWriteSplittingPlugin.SetWriterConnection", writer_host_info.url)

    def _set_reader_connection(self, reader_conn: Connection, reader_host_info: HostInfo):
        self._reader_connection = reader_conn
        self._reader_host_info = reader_host_info
        logger.debug("ReadWriteSplittingPlugin.SetReaderConnection", reader_host_info.url)

    def _get_new_writer_connection(self, writer_host: HostInfo):
        conn = self._plugin_service.connect(writer_host, self._properties)
        provider = self._conn_provider_manager.get_connection_provider(writer_host, self._properties)
        self._is_writer_conn_from_internal_pool = (ReadWriteSplittingPlugin._POOL_PROVIDER_CLASS_NAME in str(type(provider)))
        self._set_writer_connection(conn, writer_host)
        self._switch_current_connection_to(conn, writer_host)

    def _switch_connection_if_required(self, read_only: bool):
        current_conn = self._plugin_service.current_connection
        driver_dialect = self._plugin_service.driver_dialect

        if (current_conn is not None and
                driver_dialect is not None and driver_dialect.is_closed(current_conn)):
            self._log_and_raise_exception("ReadWriteSplittingPlugin.SetReadOnlyOnClosedConnection")

        if current_conn is not None and driver_dialect.can_execute_query(current_conn):
            try:
                self._plugin_service.refresh_host_list()
            except Exception:
                pass  # Swallow exception

        hosts = self._plugin_service.hosts
        if hosts is None or len(hosts) == 0:
            self._log_and_raise_exception("ReadWriteSplittingPlugin.EmptyHostList")

        current_host = self._plugin_service.current_host_info
        if current_host is None:
            self._log_and_raise_exception("ReadWriteSplittingPlugin.UnavailableHostInfo")
            return

        if read_only:
            if not self._plugin_service.is_in_transaction and current_host.role != HostRole.READER:
                try:
                    self._switch_to_reader_connection(hosts)
                except Exception:
                    if not self._is_connection_usable(current_conn, driver_dialect):
                        self._log_and_raise_exception("ReadWriteSplittingPlugin.ErrorSwitchingToReader")
                        return

                    logger.warning("ReadWriteSplittingPlugin.FallbackToWriter", current_host.url)
        elif current_host.role != HostRole.WRITER:
            if self._plugin_service.is_in_transaction:
                self._log_and_raise_exception("ReadWriteSplittingPlugin.SetReadOnlyFalseInTransaction")

            try:
                self._switch_to_writer_connection(hosts)
            except Exception:
                self._log_and_raise_exception("ReadWriteSplittingPlugin.ErrorSwitchingToWriter")

    def _switch_current_connection_to(self, new_conn: Connection, new_conn_host: HostInfo):
        current_conn = self._plugin_service.current_connection
        if current_conn == new_conn:
            return

        self._plugin_service.set_current_connection(new_conn, new_conn_host)

        logger.debug("ReadWriteSplittingPlugin.SettingCurrentConnection", new_conn_host.url)

    def _switch_to_writer_connection(self, hosts: Tuple[HostInfo, ...]):
        current_host = self._plugin_service.current_host_info
        current_conn = self._plugin_service.current_connection
        driver_dialect = self._plugin_service.driver_dialect
        if (current_host is not None and current_host.role == HostRole.WRITER and
                self._is_connection_usable(current_conn, driver_dialect)):
            return

        writer_host = self._get_writer(hosts)
        if writer_host is None:
            return

        self._in_read_write_split = True
        if not self._is_connection_usable(self._writer_connection, driver_dialect):
            self._get_new_writer_connection(writer_host)
        elif self._writer_connection is not None:
            self._switch_current_connection_to(self._writer_connection, writer_host)

        if self._is_reader_conn_from_internal_pool:
            self._close_connection_if_idle(self._reader_connection)

        logger.debug("ReadWriteSplittingPlugin.SwitchedFromReaderToWriter", writer_host.url)

    def _switch_to_reader_connection(self, hosts: Tuple[HostInfo, ...]):
        current_host = self._plugin_service.current_host_info
        current_conn = self._plugin_service.current_connection
        driver_dialect = self._plugin_service.driver_dialect
        if (current_host is not None and current_host.role == HostRole.READER and
                self._is_connection_usable(current_conn, driver_dialect)):
            return

        hostnames = [host_info.host for host_info in hosts]
        if self._reader_host_info is not None and self._reader_host_info.host not in hostnames:
            # The old reader cannot be used anymore because it is no longer in the list of allowed hosts.
            self._close_connection_if_idle(self._reader_connection)

        self._in_read_write_split = True
        if not self._is_connection_usable(self._reader_connection, driver_dialect):
            self._initialize_reader_connection(hosts)
        elif self._reader_connection is not None and self._reader_host_info is not None:
            try:
                self._switch_current_connection_to(self._reader_connection, self._reader_host_info)
                logger.debug("ReadWriteSplittingPlugin.SwitchedFromWriterToReader", self._reader_host_info.url)
            except Exception:
                logger.debug("ReadWriteSplittingPlugin.ErrorSwitchingToCachedReader", self._reader_host_info.url)

                self._reader_connection.close()
                self._reader_connection = None
                self._reader_host_info = None
                self._initialize_reader_connection(hosts)

        if self._is_writer_conn_from_internal_pool:
            self._close_connection_if_idle(self._writer_connection)

    def _initialize_reader_connection(self, hosts: Tuple[HostInfo, ...]):
        if len(hosts) == 1:
            writer_host = self._get_writer(hosts)
            if writer_host is not None:
                if not self._is_connection_usable(self._writer_connection, self._plugin_service.driver_dialect):
                    self._get_new_writer_connection(writer_host)
                logger.warning("ReadWriteSplittingPlugin.NoReadersFound", writer_host.url)
                return

        conn: Optional[Connection] = None
        reader_host: Optional[HostInfo] = None

        conn_attempts = len(self._plugin_service.hosts) * 2
        for _ in range(conn_attempts):
            host = self._plugin_service.get_host_info_by_strategy(HostRole.READER, self._reader_selector_strategy)
            if host is not None:
                try:
                    conn = self._plugin_service.connect(host, self._properties)
                    provider = self._conn_provider_manager.get_connection_provider(host, self._properties)
                    self._is_reader_conn_from_internal_pool = (ReadWriteSplittingPlugin._POOL_PROVIDER_CLASS_NAME in str(type(provider)))
                    reader_host = host
                    break
                except Exception:
                    logger.warning("ReadWriteSplittingPlugin.FailedToConnectToReader", host.url)

        if conn is None or reader_host is None:
            self._log_and_raise_exception("ReadWriteSplittingPlugin.NoReadersAvailable")
            return

        logger.debug("ReadWriteSplittingPlugin.SuccessfullyConnectedToReader", reader_host.url)

        self._set_reader_connection(conn, reader_host)
        self._switch_current_connection_to(conn, reader_host)

        logger.debug("ReadWriteSplittingPlugin.SwitchedFromWriterToReader", reader_host.url)

    def _close_connection_if_idle(self, internal_conn: Optional[Connection]):
        current_conn = self._plugin_service.current_connection
        driver_dialect = self._plugin_service.driver_dialect
        try:
            if (internal_conn is not None and internal_conn != current_conn and
                    self._is_connection_usable(internal_conn, driver_dialect)):
                internal_conn.close()
                if internal_conn == self._writer_connection:
                    self._writer_connection = None
                if internal_conn == self._reader_connection:
                    self._reader_connection = None
                    self._reader_host_info = None

        except Exception:
            pass  # Swallow exception

    def _close_idle_connections(self):
        logger.debug("ReadWriteSplittingPlugin.ClosingInternalConnections")
        self._close_connection_if_idle(self._reader_connection)
        self._close_connection_if_idle(self._writer_connection)

    @staticmethod
    def _log_and_raise_exception(log_msg: str):
        logger.error(log_msg)
        raise ReadWriteSplittingError(Messages.get(log_msg))

    @staticmethod
    def _is_connection_usable(conn: Optional[Connection], driver_dialect: Optional[DriverDialect]):
        return conn is not None and driver_dialect is not None and not driver_dialect.is_closed(conn)

    @staticmethod
    def _get_writer(hosts: Tuple[HostInfo, ...]) -> Optional[HostInfo]:
        for host in hosts:
            if host.role == HostRole.WRITER:
                return host

        ReadWriteSplittingPlugin._log_and_raise_exception("ReadWriteSplittingPlugin.NoWriterFound")

        return None


class ReadWriteSplittingPluginFactory(PluginFactory):
    def get_instance(self, plugin_service: PluginService, props: Properties) -> Plugin:
        return ReadWriteSplittingPlugin(plugin_service, props)