custom_endpoint_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 threading import Event, Thread
from time import perf_counter_ns, sleep
from typing import (TYPE_CHECKING, Any, Callable, ClassVar, Dict, List,
                    Optional, Set, Union, cast)

from aws_advanced_python_wrapper.allowed_and_blocked_hosts import \
    AllowedAndBlockedHosts
from aws_advanced_python_wrapper.errors import AwsWrapperError
from aws_advanced_python_wrapper.utils.cache_map import CacheMap
from aws_advanced_python_wrapper.utils.messages import Messages
from aws_advanced_python_wrapper.utils.region_utils import RegionUtils

if TYPE_CHECKING:
    from aws_advanced_python_wrapper.driver_dialect import DriverDialect
    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 enum import Enum

from boto3 import Session  # type: ignore

from aws_advanced_python_wrapper.pep249_methods import DbApiMethod
from aws_advanced_python_wrapper.plugin import Plugin, PluginFactory
from aws_advanced_python_wrapper.utils.log import Logger
from aws_advanced_python_wrapper.utils.properties import WrapperProperties
from aws_advanced_python_wrapper.utils.rds_utils import RdsUtils
from aws_advanced_python_wrapper.utils.sliding_expiration_cache_container import \
    SlidingExpirationCacheContainer
from aws_advanced_python_wrapper.utils.telemetry.telemetry import (
    TelemetryCounter, TelemetryFactory)

logger = Logger(__name__)


class CustomEndpointRoleType(Enum):
    """
    Enum representing the possible roles of instances specified by a custom endpoint. Note that, currently, it is not
    possible to create a WRITER custom endpoint.
    """
    ANY = "ANY"
    READER = "READER"

    @classmethod
    def from_string(cls, value):
        return CustomEndpointRoleType(value)


class CustomEndpointInfo:
    def __init__(self,
                 endpoint_id: str,
                 cluster_id: str,
                 endpoint: str,
                 role_type: CustomEndpointRoleType,
                 static_members: Optional[Set[str]],
                 excluded_members: Optional[Set[str]]):
        self.endpoint_id = endpoint_id
        self.cluster_id = cluster_id
        self.endpoint = endpoint
        self.role_type = role_type
        self.static_members = None if not static_members else static_members
        self.excluded_members = None if not excluded_members else excluded_members

    @classmethod
    def from_db_cluster_endpoint(cls, endpoint_response_info: Dict[str, Union[str, List[str]]]):
        return CustomEndpointInfo(
            str(endpoint_response_info.get("DBClusterEndpointIdentifier")),
            str(endpoint_response_info.get("DBClusterIdentifier")),
            str(endpoint_response_info.get("Endpoint")),
            CustomEndpointRoleType.from_string(str(endpoint_response_info.get("CustomEndpointType"))),
            set(cast('List[str]', endpoint_response_info.get("StaticMembers"))),
            set(cast('List[str]', endpoint_response_info.get("ExcludedMembers")))
        )

    def __eq__(self, other: object):
        if self is object:
            return True
        if not isinstance(other, CustomEndpointInfo):
            return False

        return self.endpoint_id == other.endpoint_id \
            and self.cluster_id == other.cluster_id \
            and self.endpoint == other.endpoint \
            and self.role_type == other.role_type \
            and self.static_members == other.static_members \
            and self.excluded_members == other.excluded_members

    def __hash__(self):
        return hash((self.endpoint_id, self.cluster_id, self.endpoint, self.role_type))

    def __str__(self):
        return (f"CustomEndpointInfo[endpoint={self.endpoint}, cluster_id={self.cluster_id}, "
                f"role_type={self.role_type}, endpoint_id={self.endpoint_id}, static_members={self.static_members}, "
                f"excluded_members={self.excluded_members}]")


class CustomEndpointMonitor:
    """
    A custom endpoint monitor. This class uses a background thread to monitor a given custom endpoint for custom
    endpoint information and future changes to the custom endpoint.
    """
    _CUSTOM_ENDPOINT_INFO_EXPIRATION_NS: ClassVar[int] = 5 * 60_000_000_000  # 5 minutes
    # Keys are custom endpoint URLs, values are information objects for the associated custom endpoint.
    _custom_endpoint_info_cache: ClassVar[CacheMap[str, CustomEndpointInfo]] = CacheMap()

    def __init__(self,
                 plugin_service: PluginService,
                 custom_endpoint_host_info: HostInfo,
                 endpoint_id: str,
                 region: str,
                 refresh_rate_ns: int,
                 session: Optional[Session] = None):
        self._plugin_service = plugin_service
        self._custom_endpoint_host_info = custom_endpoint_host_info
        self._endpoint_id = endpoint_id
        self._region = region
        self._refresh_rate_ns = refresh_rate_ns
        self._session = session if session else Session()
        self._client = self._session.client('rds', region_name=region)

        self._stop_event = Event()
        telemetry_factory = self._plugin_service.get_telemetry_factory()
        self._info_changed_counter = telemetry_factory.create_counter("customEndpoint.infoChanged.counter")

        self._thread = Thread(daemon=True, name="CustomEndpointMonitorThread", target=self._run)
        self._thread.start()

    def _run(self):
        logger.debug("CustomEndpointMonitor.StartingMonitor", self._custom_endpoint_host_info.host)

        try:
            while not self._stop_event.is_set():
                try:
                    start_ns = perf_counter_ns()

                    response = self._client.describe_db_cluster_endpoints(
                        DBClusterEndpointIdentifier=self._endpoint_id,
                        Filters=[
                            {
                                "Name": "db-cluster-endpoint-type",
                                "Values": ["custom"]
                            }
                        ]
                    )

                    endpoints = response["DBClusterEndpoints"]
                    if len(endpoints) != 1:
                        endpoint_hostnames = [endpoint["Endpoint"] for endpoint in endpoints]
                        logger.warning(
                            "CustomEndpointMonitor.UnexpectedNumberOfEndpoints",
                            self._endpoint_id,
                            self._region,
                            len(endpoints),
                            endpoint_hostnames)

                        if self._stop_event.wait(self._refresh_rate_ns / 1_000_000_000):
                            break
                        continue

                    endpoint_info = CustomEndpointInfo.from_db_cluster_endpoint(endpoints[0])
                    cached_info = \
                        CustomEndpointMonitor._custom_endpoint_info_cache.get(self._custom_endpoint_host_info.host)
                    if cached_info is not None and cached_info == endpoint_info:
                        elapsed_time = perf_counter_ns() - start_ns
                        sleep_duration = max(0, self._refresh_rate_ns - elapsed_time)
                        if self._stop_event.wait(sleep_duration / 1_000_000_000):
                            break
                        continue

                    logger.debug(
                        "CustomEndpointMonitor.DetectedChangeInCustomEndpointInfo",
                        self._custom_endpoint_host_info.host, endpoint_info)

                    # The custom endpoint info has changed, so we need to update the set of allowed/blocked hosts.
                    hosts = AllowedAndBlockedHosts(endpoint_info.static_members, endpoint_info.excluded_members)
                    self._plugin_service.allowed_and_blocked_hosts = hosts
                    CustomEndpointMonitor._custom_endpoint_info_cache.put(
                        self._custom_endpoint_host_info.host,
                        endpoint_info,
                        CustomEndpointMonitor._CUSTOM_ENDPOINT_INFO_EXPIRATION_NS)

                    if self._info_changed_counter is not None:
                        self._info_changed_counter.inc()

                    elapsed_time = perf_counter_ns() - start_ns
                    sleep_duration = max(0, self._refresh_rate_ns - elapsed_time)
                    if self._stop_event.wait(sleep_duration / 1_000_000_000):
                        break
                    continue
                except InterruptedError as e:
                    raise e
                except Exception as e:
                    # If the exception is not an InterruptedError, log it and continue monitoring.
                    logger.error("CustomEndpointMonitor.Exception", self._custom_endpoint_host_info.host, e)
        except InterruptedError:
            logger.info("CustomEndpointMonitor.Interrupted", self._custom_endpoint_host_info.host)
        finally:
            CustomEndpointMonitor._custom_endpoint_info_cache.remove(self._custom_endpoint_host_info.host)
            self._stop_event.set()
            self._client.close()
            logger.debug("CustomEndpointMonitor.StoppedMonitor", self._custom_endpoint_host_info.host)

    def has_custom_endpoint_info(self):
        return CustomEndpointMonitor._custom_endpoint_info_cache.get(self._custom_endpoint_host_info.host) is not None

    def close(self):
        logger.debug("CustomEndpointMonitor.StoppingMonitor", self._custom_endpoint_host_info.host)
        CustomEndpointMonitor._custom_endpoint_info_cache.remove(self._custom_endpoint_host_info.host)
        self._stop_event.set()


class CustomEndpointPlugin(Plugin):
    """
    A plugin that analyzes custom endpoints for custom endpoint information and custom endpoint changes, such as adding
    or removing an instance in the custom endpoint.
    """
    _SUBSCRIBED_METHODS: ClassVar[Set[str]] = {DbApiMethod.CONNECT.method_name}
    _CACHE_CLEANUP_RATE_NS: ClassVar[int] = 60_000_000_000  # 1 minute
    _MONITOR_CACHE_NAME: ClassVar[str] = "custom_endpoint_monitors"

    def __init__(self, plugin_service: PluginService, props: Properties):
        self._plugin_service = plugin_service
        self._props = props

        self._should_wait_for_info: bool = WrapperProperties.WAIT_FOR_CUSTOM_ENDPOINT_INFO.get_bool(self._props)
        self._wait_for_info_timeout_ms: int = WrapperProperties.WAIT_FOR_CUSTOM_ENDPOINT_INFO_TIMEOUT_MS.get_int(self._props)
        self._idle_monitor_expiration_ms: int = \
            WrapperProperties.CUSTOM_ENDPOINT_IDLE_MONITOR_EXPIRATION_MS.get_int(self._props)

        self._rds_utils = RdsUtils()
        self._region_utils = RegionUtils()
        self._region: Optional[str] = None
        self._custom_endpoint_host_info: Optional[HostInfo] = None
        self._custom_endpoint_id: Optional[str] = None
        telemetry_factory: TelemetryFactory = self._plugin_service.get_telemetry_factory()
        self._wait_for_info_counter: TelemetryCounter | None = telemetry_factory.create_counter("customEndpoint.waitForInfo.counter")

        self._monitors = SlidingExpirationCacheContainer.get_or_create_cache(
            name=CustomEndpointPlugin._MONITOR_CACHE_NAME,
            cleanup_interval_ns=CustomEndpointPlugin._CACHE_CLEANUP_RATE_NS,
            should_dispose_func=lambda _: True,
            item_disposal_func=lambda monitor: monitor.close()
        )

        CustomEndpointPlugin._SUBSCRIBED_METHODS.update(self._plugin_service.network_bound_methods)

    @property
    def subscribed_methods(self) -> Set[str]:
        return CustomEndpointPlugin._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:
        if not self._rds_utils.is_rds_custom_cluster_dns(host_info.host):
            return connect_func()

        self._custom_endpoint_host_info = host_info
        logger.debug("CustomEndpointPlugin.ConnectionRequestToCustomEndpoint", host_info.host)

        self._custom_endpoint_id = self._rds_utils.get_cluster_id(host_info.host)
        if not self._custom_endpoint_id:
            raise AwsWrapperError(
                Messages.get_formatted(
                    "CustomEndpointPlugin.ErrorParsingEndpointIdentifier", self._custom_endpoint_host_info.host))

        hostname = self._custom_endpoint_host_info.host
        self._region = self._region_utils.get_region_from_hostname(hostname)
        if not self._region:
            error_message = "RdsUtils.UnsupportedHostname"
            logger.debug(error_message, hostname)
            raise AwsWrapperError(Messages.get_formatted(error_message, hostname))

        monitor = self._create_monitor_if_absent(props)
        if self._should_wait_for_info:
            self._wait_for_info(monitor)

        return connect_func()

    def _create_monitor_if_absent(self, props: Properties) -> CustomEndpointMonitor:
        host_info = cast('HostInfo', self._custom_endpoint_host_info)
        endpoint_id = cast('str', self._custom_endpoint_id)
        region = cast('str', self._region)
        monitor = self._monitors.compute_if_absent(
            host_info.host,
            lambda key: CustomEndpointMonitor(
                self._plugin_service,
                host_info,
                endpoint_id,
                region,
                WrapperProperties.CUSTOM_ENDPOINT_INFO_REFRESH_RATE_MS.get_int(props) * 1_000_000),
            self._idle_monitor_expiration_ms * 1_000_000)

        return cast('CustomEndpointMonitor', monitor)

    def _wait_for_info(self, monitor: CustomEndpointMonitor):
        has_info = monitor.has_custom_endpoint_info()
        if has_info:
            return

        if self._wait_for_info_counter is not None:
            self._wait_for_info_counter.inc()
        host_info = cast('HostInfo', self._custom_endpoint_host_info)
        hostname = host_info.host
        logger.debug("CustomEndpointPlugin.WaitingForCustomEndpointInfo", hostname, self._wait_for_info_timeout_ms)
        wait_for_info_timeout_ns = perf_counter_ns() + self._wait_for_info_timeout_ms * 1_000_000

        try:
            while not has_info and perf_counter_ns() < wait_for_info_timeout_ns:
                sleep(0.1)
                has_info = monitor.has_custom_endpoint_info()
        except InterruptedError as e:
            raise AwsWrapperError(Messages.get_formatted("CustomEndpointPlugin.InterruptedThread", hostname), e) from e

        if not has_info:
            raise AwsWrapperError(
                Messages.get_formatted(
                    "CustomEndpointPlugin.TimedOutWaitingForCustomEndpointInfo",
                    self._wait_for_info_timeout_ms, hostname))

    def execute(self, target: type, method_name: str, execute_func: Callable, *args: Any, **kwargs: Any) -> Any:
        if self._custom_endpoint_host_info is None:
            return execute_func()

        monitor = self._create_monitor_if_absent(self._props)
        if self._should_wait_for_info:
            self._wait_for_info(monitor)

        return execute_func()


class CustomEndpointPluginFactory(PluginFactory):
    @staticmethod
    def get_instance(plugin_service: PluginService, props: Properties) -> Plugin:
        return CustomEndpointPlugin(plugin_service, props)