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

aws_advanced_python_wrapper/host_selector.py

@@ -0,0 +1,190 @@
+#  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 random
+from re import search
+from typing import TYPE_CHECKING, Dict, List, Optional, Protocol, Tuple
+
+from .host_availability import HostAvailability
+
+if TYPE_CHECKING:
+    from .hostinfo import HostInfo, HostRole
+
+from aws_advanced_python_wrapper.errors import AwsWrapperError
+from aws_advanced_python_wrapper.utils.cache_map import CacheMap
+from .pep249 import Error
+from .utils.messages import Messages
+from .utils.properties import Properties, WrapperProperties
+
+
+class HostSelector(Protocol):
+    """
+    Interface for a strategy defining how to pick a host from a list of hosts.
+    """
+
+    def get_host(self, hosts: Tuple[HostInfo, ...], role: HostRole, props: Optional[Properties] = None) -> HostInfo:
+        ...
+
+
+class RandomHostSelector(HostSelector):
+
+    def get_host(self, hosts: Tuple[HostInfo, ...], role: HostRole, props: Optional[Properties] = None) -> HostInfo:
+
+        eligible_hosts = [host for host in hosts if host.role == role and host.get_availability() == HostAvailability.AVAILABLE]
+
+        if len(eligible_hosts) == 0:
+            raise Error(Messages.get("HostSelector.NoEligibleHost"))
+
+        return random.choice(eligible_hosts)
+
+
+class RoundRobinClusterInfo:
+    _last_host: Optional[HostInfo] = None
+    _cluster_weights_dict: Dict[str, int] = {}
+    _default_weight: int = 1
+    _weight_counter: int = 0
+
+    @property
+    def last_host(self) -> Optional[HostInfo]:
+        return self._last_host
+
+    @last_host.setter
+    def last_host(self, value):
+        self._last_host = value
+
+    @property
+    def cluster_weights_dict(self) -> Dict[str, int]:
+        return self._cluster_weights_dict
+
+    @cluster_weights_dict.setter
+    def cluster_weights_dict(self, value):
+        self._cluster_weights_dict = value
+
+    @property
+    def default_weight(self):
+        return self._default_weight
+
+    @default_weight.setter
+    def default_weight(self, value):
+        self._default_weight = value
+
+    @property
+    def weight_counter(self) -> int:
+        return self._weight_counter
+
+    @weight_counter.setter
+    def weight_counter(self, value):
+        self._weight_counter = value
+
+
+class RoundRobinHostSelector(HostSelector):
+    _DEFAULT_WEIGHT: int = 1
+    _DEFAULT_ROUND_ROBIN_CACHE_EXPIRE_NANOS = 60000000000 * 10  # 10 minutes
+    _HOST_WEIGHT_PAIRS_PATTERN = r"((?P<host>[^:/?#]*):(?P<weight>.*))"
+    _round_robin_cache: CacheMap[str, Optional[RoundRobinClusterInfo]] = CacheMap()
+
+    def get_host(self, hosts: Tuple[HostInfo, ...], role: HostRole, props: Optional[Properties] = None) -> HostInfo:
+
+        eligible_hosts: List[HostInfo] = [host for host in hosts if host.role == role and host.get_availability() == HostAvailability.AVAILABLE]
+        eligible_hosts.sort(key=lambda host: host.host, reverse=False)
+        if len(eligible_hosts) == 0:
+            raise AwsWrapperError(Messages.get_formatted("HostSelector.NoHostsMatchingRole", role))
+
+        # Create new cache entries for provided hosts if necessary. All hosts point to the same cluster info.
+        self._create_cache_entry_for_hosts(eligible_hosts, props)
+        current_cluster_info_key: str = eligible_hosts[0].host
+        cluster_info: Optional[RoundRobinClusterInfo] = RoundRobinHostSelector._round_robin_cache.get(current_cluster_info_key)
+
+        last_host_index: int = -1
+        if cluster_info is None:
+            raise AwsWrapperError(Messages.get("RoundRobinHostSelector.ClusterInfoNone"))
+
+        last_host = cluster_info.last_host
+        # Check if last_host is in list of eligible hosts. Update last_host_index.
+        if last_host is not None:
+            for i in range(0, len(eligible_hosts)):
+                if eligible_hosts[i].host == last_host.host:
+                    last_host_index = i
+
+        if cluster_info.weight_counter > 0 and last_host_index != -1:
+            target_host_index = last_host_index
+        else:
+            if last_host_index != -1 and last_host_index != (len(eligible_hosts) - 1):
+                target_host_index = last_host_index + 1
+            else:
+                target_host_index = 0
+            weight = cluster_info.cluster_weights_dict.get(eligible_hosts[target_host_index].host)
+            cluster_info.weight_counter = cluster_info.default_weight if weight is None else weight
+
+        cluster_info.weight_counter = (cluster_info.weight_counter - 1)
+        cluster_info.last_host = eligible_hosts[target_host_index]
+        return eligible_hosts[target_host_index]
+
+    def _create_cache_entry_for_hosts(self, hosts: List[HostInfo], props: Optional[Properties]) -> None:
+        cached_info = None
+        for host in hosts:
+            info = self._round_robin_cache.get(host.host)
+            if info is not None:
+                cached_info = info
+                break
+        if cached_info is not None:
+            for host in hosts:
+                # Update the expiration time
+                self._round_robin_cache.put(
+                    host.host, cached_info, RoundRobinHostSelector._DEFAULT_ROUND_ROBIN_CACHE_EXPIRE_NANOS)
+        else:
+            round_robin_cluster_info: RoundRobinClusterInfo = RoundRobinClusterInfo()
+            self._update_cache_properties_for_round_robin_cluster_info(round_robin_cluster_info, props)
+            for host in hosts:
+                self._round_robin_cache.put(
+                    host.host, round_robin_cluster_info, RoundRobinHostSelector._DEFAULT_ROUND_ROBIN_CACHE_EXPIRE_NANOS)
+
+    def _update_cache_properties_for_round_robin_cluster_info(self, round_robin_cluster_info: RoundRobinClusterInfo, props: Optional[Properties]):
+        cluster_default_weight: int = RoundRobinHostSelector._DEFAULT_WEIGHT
+        if props is not None:
+            props_weight = WrapperProperties.ROUND_ROBIN_DEFAULT_WEIGHT.get_int(props)
+            if props_weight < RoundRobinHostSelector._DEFAULT_WEIGHT:
+                raise AwsWrapperError(Messages.get("RoundRobinHostSelector.RoundRobinInvalidDefaultWeight"))
+            cluster_default_weight = props_weight
+        round_robin_cluster_info.default_weight = cluster_default_weight
+
+        if props is not None:
+            host_weights: Optional[str] = WrapperProperties.ROUND_ROBIN_HOST_WEIGHT_PAIRS.get(props)
+            if host_weights is not None and len(host_weights) != 0:
+                host_weight_pairs: List[str] = host_weights.split(",")
+
+                for pair in host_weight_pairs:
+                    match = search(RoundRobinHostSelector._HOST_WEIGHT_PAIRS_PATTERN, pair)
+                    if match:
+                        host_name = match.group("host")
+                        host_weight = match.group("weight")
+                    else:
+                        raise AwsWrapperError(Messages.get("RoundRobinHostSelector.RoundRobinInvalidHostWeightPairs"))
+
+                    if len(host_name) == 0 or len(host_weight) == 0:
+                        raise AwsWrapperError(Messages.get("RoundRobinHostSelector.RoundRobinInvalidHostWeightPairs"))
+                    try:
+                        weight: int = int(host_weight)
+
+                        if weight < RoundRobinHostSelector._DEFAULT_WEIGHT:
+                            raise AwsWrapperError(Messages.get("RoundRobinHostSelector.RoundRobinInvalidHostWeightPairs"))
+
+                        round_robin_cluster_info.cluster_weights_dict[host_name] = weight
+                    except ValueError:
+                        raise AwsWrapperError(Messages.get("RoundRobinHostSelector.RoundRobinInvalidHostWeightPairs"))
+
+    def clear_cache(self):
+        RoundRobinHostSelector._round_robin_cache.clear()

aws_advanced_python_wrapper/hostinfo.py

@@ -0,0 +1,138 @@
+#  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 dataclasses import dataclass
+from enum import Enum, auto
+from typing import TYPE_CHECKING, ClassVar, FrozenSet, Optional, Set
+
+from aws_advanced_python_wrapper.host_availability import (
+    HostAvailability, HostAvailabilityStrategy)
+
+if TYPE_CHECKING:
+    from datetime import datetime
+
+
+class HostRole(Enum):
+    UNKNOWN = auto()
+    READER = auto()
+    WRITER = auto()
+
+
+@dataclass(eq=False)
+class HostInfo:
+    NO_PORT: ClassVar[int] = -1
+    DEFAULT_WEIGHT = 100
+
+    def __init__(
+            self,
+            host: str,
+            port: int = NO_PORT,
+            role: HostRole = HostRole.WRITER,
+            availability: HostAvailability = HostAvailability.AVAILABLE,
+            host_availability_strategy=HostAvailabilityStrategy(),
+            weight: int = DEFAULT_WEIGHT,
+            host_id: Optional[str] = None,
+            last_update_time: Optional[datetime] = None):
+        self.host = host
+        self.port = port
+        self.role = role
+        self._availability = availability
+        self.host_availability_strategy = host_availability_strategy
+        self.weight = weight,
+        self.host_id = host_id
+        self.last_update_time = last_update_time
+
+        self._aliases: Set[str] = set()
+        self._all_aliases: Set[str] = {self.as_alias()}
+
+    def __eq__(self, other: object):
+        if self is object:
+            return True
+        if not isinstance(other, HostInfo):
+            return False
+
+        return self.host == other.host \
+            and self.port == other.port \
+            and self._availability == other._availability \
+            and self.role == other.role
+
+    def __str__(self):
+        return f"HostInfo({self.host}, {self.port}, {self.role}, {self._availability})"
+
+    @property
+    def url(self):
+        if self.is_port_specified():
+            return f"{self.host}:{self.port}"
+        else:
+            return self.host
+
+    @property
+    def aliases(self) -> FrozenSet[str]:
+        return frozenset(self._aliases)
+
+    @property
+    def all_aliases(self) -> FrozenSet[str]:
+        return frozenset(self._all_aliases)
+
+    def as_alias(self) -> str:
+        return f"{self.host}:{self.port}" if self.is_port_specified() else self.host
+
+    def add_alias(self, *aliases: str):
+        if not aliases:
+            return
+
+        for alias in aliases:
+            self._aliases.add(alias)
+            self._all_aliases.add(alias)
+
+    def as_aliases(self) -> FrozenSet[str]:
+        return frozenset(self.all_aliases)
+
+    def remove_alias(self, *kwargs):
+        if not kwargs or len(kwargs) == 0:
+            return
+
+        for x in kwargs:
+            self._aliases.discard(x)
+            self._all_aliases.discard(x)
+
+    def reset_aliases(self):
+        self._aliases.clear()
+        self._all_aliases.clear()
+        self._all_aliases.add(self.as_alias())
+
+    def is_port_specified(self) -> bool:
+        return self.port != HostInfo.NO_PORT
+
+    def get_availability(self) -> HostAvailability:
+        if self.host_availability_strategy is not None:
+            return self.host_availability_strategy.get_host_availability(self._availability)
+
+        return self._availability
+
+    def get_raw_availability(self) -> HostAvailability:
+        return self._availability
+
+    def set_availability(self, availability: HostAvailability):
+        self._availability = availability
+        if self.host_availability_strategy is not None:
+            self.host_availability_strategy.set_host_availability(availability)
+
+    def get_host_availability_strategy(self):
+        return self.host_availability_strategy
+
+    def set_host_availability_strategy(self, host_availability_strategy):
+        self.host_availability_strategy = host_availability_strategy

aws_advanced_python_wrapper/iam_plugin.py

@@ -0,0 +1,195 @@
+#  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
+
+from aws_advanced_python_wrapper.utils.iamutils import IamAuthUtils, TokenInfo
+
+if TYPE_CHECKING:
+    from boto3 import Session
+    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 datetime import datetime, timedelta
+from typing import Callable, Dict, Optional, Set
+
+import boto3
+
+from aws_advanced_python_wrapper.errors import AwsWrapperError
+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.properties import (Properties,
+                                                          WrapperProperties)
+from aws_advanced_python_wrapper.utils.rdsutils import RdsUtils
+from aws_advanced_python_wrapper.utils.telemetry.telemetry import \
+    TelemetryTraceLevel
+
+logger = Logger(__name__)
+
+
+class IamAuthPlugin(Plugin):
+    _SUBSCRIBED_METHODS: Set[str] = {"connect", "force_connect"}
+    # Leave 30 second buffer to prevent time-of-check to time-of-use errors
+    _DEFAULT_TOKEN_EXPIRATION_SEC = 15 * 60 - 30
+
+    _rds_utils: RdsUtils = RdsUtils()
+    _token_cache: Dict[str, TokenInfo] = {}
+
+    def __init__(self, plugin_service: PluginService, session: Optional[Session] = None):
+        self._plugin_service = plugin_service
+        self._session = session
+
+        telemetry_factory = self._plugin_service.get_telemetry_factory()
+        self._fetch_token_counter = telemetry_factory.create_counter("iam.fetch_token.count")
+        self._cache_size_gauge = telemetry_factory.create_gauge(
+            "iam.token_cache.size", lambda: len(IamAuthPlugin._token_cache))
+
+    @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._connect(host_info, props, connect_func)
+
+    def _connect(self, host_info: HostInfo, props: Properties, connect_func: Callable) -> Connection:
+        if not WrapperProperties.USER.get(props):
+            raise AwsWrapperError(Messages.get_formatted("IamPlugin.IsNoneOrEmpty", WrapperProperties.USER.name))
+
+        host = IamAuthUtils.get_iam_host(props, host_info)
+        region = WrapperProperties.IAM_REGION.get(props) \
+            if WrapperProperties.IAM_REGION.get(props) else self._get_rds_region(host)
+        port = IamAuthUtils.get_port(props, host_info, self._plugin_service.database_dialect.default_port)
+        token_expiration_sec: int = WrapperProperties.IAM_EXPIRATION.get_int(props)
+
+        cache_key: str = self._get_cache_key(
+            WrapperProperties.USER.get(props),
+            host,
+            port,
+            region
+        )
+
+        token_info = IamAuthPlugin._token_cache.get(cache_key)
+
+        if token_info is not None and not token_info.is_expired():
+            logger.debug("IamAuthPlugin.UseCachedIamToken", token_info.token)
+            self._plugin_service.driver_dialect.set_password(props, token_info.token)
+        else:
+            token_expiry = datetime.now() + timedelta(seconds=token_expiration_sec)
+            token: str = self._generate_authentication_token(props, host, port, region)
+            logger.debug("IamAuthPlugin.GeneratedNewIamToken", token)
+            self._plugin_service.driver_dialect.set_password(props, token)
+            IamAuthPlugin._token_cache[cache_key] = TokenInfo(token, token_expiry)
+
+        try:
+            return connect_func()
+
+        except Exception as e:
+            logger.debug("IamAuthPlugin.ConnectException", e)
+
+            is_cached_token = (token_info is not None and not token_info.is_expired())
+            if not self._plugin_service.is_login_exception(error=e) or not is_cached_token:
+                raise AwsWrapperError(Messages.get_formatted("IamAuthPlugin.ConnectException", e)) from e
+
+            # Login unsuccessful with cached token
+            # Try to generate a new token and try to connect again
+            token_expiry = datetime.now() + timedelta(seconds=token_expiration_sec)
+            token = self._generate_authentication_token(props, host, port, region)
+            logger.debug("IamAuthPlugin.GeneratedNewIamToken", token)
+            self._plugin_service.driver_dialect.set_password(props, token)
+            IamAuthPlugin._token_cache[token] = TokenInfo(token, token_expiry)
+
+            try:
+                return connect_func()
+            except Exception as e:
+                raise AwsWrapperError(Messages.get_formatted("IamAuthPlugin.UnhandledException", e)) from e
+
+    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, force_connect_func)
+
+    def _generate_authentication_token(self,
+                                       props: Properties,
+                                       hostname: Optional[str],
+                                       port: Optional[int],
+                                       region: Optional[str]) -> str:
+        telemetry_factory = self._plugin_service.get_telemetry_factory()
+        context = telemetry_factory.open_telemetry_context("fetch IAM token", TelemetryTraceLevel.NESTED)
+        self._fetch_token_counter.inc()
+
+        try:
+            session = self._session if self._session else boto3.Session()
+            client = session.client(
+                'rds',
+                region_name=region,
+            )
+
+            user = WrapperProperties.USER.get(props)
+
+            token = client.generate_db_auth_token(
+                DBHostname=hostname,
+                Port=port,
+                DBUsername=user
+            )
+
+            client.close()
+
+            return token
+        except Exception as ex:
+            context.set_success(False)
+            context.set_exception(ex)
+            raise ex
+        finally:
+            context.close_context()
+
+    def _get_cache_key(self, user: Optional[str], hostname: Optional[str], port: int, region: Optional[str]) -> str:
+        return f"{region}:{hostname}:{port}:{user}"
+
+    def _get_rds_region(self, hostname: Optional[str]) -> str:
+        rds_region = self._rds_utils.get_rds_region(hostname) if hostname else None
+
+        if not rds_region:
+            exception_message = "RdsUtils.UnsupportedHostname"
+            logger.debug(exception_message, hostname)
+            raise AwsWrapperError(Messages.get_formatted(exception_message, hostname))
+
+        session = self._session if self._session else boto3.Session()
+        if rds_region not in session.get_available_regions("rds"):
+            exception_message = "AwsSdk.UnsupportedRegion"
+            logger.debug(exception_message, rds_region)
+            raise AwsWrapperError(Messages.get_formatted(exception_message, rds_region))
+
+        return rds_region
+
+
+class IamAuthPluginFactory(PluginFactory):
+    def get_instance(self, plugin_service: PluginService, props: Properties) -> Plugin:
+        return IamAuthPlugin(plugin_service)