feat: limitless plugin implementation

Author: JuanLeee

aws_advanced_python_wrapper/database_dialect.py

@@ -98,6 +98,15 @@ def is_reader_query(self) -> str:
         return self._IS_READER_QUERY
 
 
+@runtime_checkable
+class AuroraLimitlessDialect(Protocol):
+    _LIMITLESS_ROUTER_ENDPOINT_QUERY: str
+
+    @property
+    def limitless_router_endpoint_query(self) -> str:
+        return self._LIMITLESS_ROUTER_ENDPOINT_QUERY
+
+
 class DatabaseDialect(Protocol):
     """
     Database dialects help the AWS Advanced Python Driver determine what kind of underlying database is being used,
@@ -342,7 +351,7 @@ def get_host_list_provider_supplier(self) -> Callable:
         return lambda provider_service, props: RdsHostListProvider(provider_service, props)
 
 
-class AuroraPgDialect(PgDatabaseDialect, TopologyAwareDatabaseDialect):
+class AuroraPgDialect(PgDatabaseDialect, TopologyAwareDatabaseDialect, AuroraLimitlessDialect):
     _DIALECT_UPDATE_CANDIDATES: Tuple[DialectCode, ...] = (DialectCode.MULTI_AZ_PG,)
 
     _EXTENSIONS_QUERY = "SELECT (setting LIKE '%aurora_stat_utils%') AS aurora_stat_utils " \
@@ -359,6 +368,7 @@ class AuroraPgDialect(PgDatabaseDialect, TopologyAwareDatabaseDialect):
 
     _HOST_ID_QUERY = "SELECT aurora_db_instance_identifier()"
     _IS_READER_QUERY = "SELECT pg_is_in_recovery()"
+    _LIMITLESS_ROUTER_ENDPOINT_QUERY = "SELECT router_endpoint, load FROM aurora_limitless_router_endpoints()"
 
     @property
     def dialect_update_candidates(self) -> Optional[Tuple[DialectCode, ...]]:

aws_advanced_python_wrapper/default_plugin.py

@@ -14,7 +14,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, List, Optional
 
 if TYPE_CHECKING:
     from aws_advanced_python_wrapper.connection_provider import (ConnectionProvider,
@@ -118,7 +118,7 @@ def accepts_strategy(self, role: HostRole, strategy: str) -> bool:
             return False
         return self._connection_provider_manager.accepts_strategy(role, strategy)
 
-    def get_host_info_by_strategy(self, role: HostRole, strategy: str) -> HostInfo:
+    def get_host_info_by_strategy(self, role: HostRole, strategy: str, host_list: Optional[List[HostInfo]] = None) -> HostInfo:
         if HostRole.UNKNOWN == role:
             raise AwsWrapperError(Messages.get("DefaultPlugin.UnknownHosts"))
 
@@ -127,7 +127,10 @@ def get_host_info_by_strategy(self, role: HostRole, strategy: str) -> HostInfo:
         if len(hosts) < 1:
             raise AwsWrapperError(Messages.get("DefaultPlugin.EmptyHosts"))
 
-        return self._connection_provider_manager.get_host_info_by_strategy(hosts, role, strategy, self._plugin_service.props)
+        if host_list is None:
+            return self._connection_provider_manager.get_host_info_by_strategy(hosts, role, strategy, self._plugin_service.props)
+        else:
+            return self._connection_provider_manager.get_host_info_by_strategy(tuple(host_list), role, strategy, self._plugin_service.props)
 
     @property
     def subscribed_methods(self) -> Set[str]:

aws_advanced_python_wrapper/fastest_response_strategy_plugin.py

@@ -58,8 +58,8 @@ def __init__(self, plugin_service: PluginService, props: Properties):
         self._plugin_service = plugin_service
         self._properties = props
         self._host_response_time_service: HostResponseTimeService = \
-            HostResponseTimeService(plugin_service, props, WrapperProperties.RESPONSE_MEASUREMENT_INTERVAL_MILLIS.get_int(props))
-        self._cache_expiration_nanos = WrapperProperties.RESPONSE_MEASUREMENT_INTERVAL_MILLIS.get_int(props) * 10 ^ 6
+            HostResponseTimeService(plugin_service, props, WrapperProperties.RESPONSE_MEASUREMENT_INTERVAL_MS.get_int(props))
+        self._cache_expiration_nanos = WrapperProperties.RESPONSE_MEASUREMENT_INTERVAL_MS.get_int(props) * 10 ^ 6
         self._random_host_selector = RandomHostSelector()
         self._cached_fastest_response_host_by_role: CacheMap[str, HostInfo] = CacheMap()
         self._hosts: Tuple[HostInfo, ...] = ()
@@ -86,7 +86,7 @@ def connect(
     def accepts_strategy(self, role: HostRole, strategy: str) -> bool:
         return strategy == FastestResponseStrategyPlugin._FASTEST_RESPONSE_STRATEGY_NAME
 
-    def get_host_info_by_strategy(self, role: HostRole, strategy: str) -> HostInfo:
+    def get_host_info_by_strategy(self, role: HostRole, strategy: str, host_list: Optional[List[HostInfo]] = None) -> HostInfo:
         if not self.accepts_strategy(role, strategy):
             logger.error("FastestResponseStrategyPlugin.UnsupportedHostSelectorStrategy", strategy)
             raise AwsWrapperError(Messages.get_formatted("FastestResponseStrategyPlugin.UnsupportedHostSelectorStrategy", strategy))

aws_advanced_python_wrapper/host_list_provider.py

@@ -69,6 +69,9 @@ def get_host_role(self, connection: Connection) -> HostRole:
     def identify_connection(self, connection: Optional[Connection]) -> Optional[HostInfo]:
         ...
 
+    def get_cluster_id(self) -> str:
+        ...
+
 
 @runtime_checkable
 class DynamicHostListProvider(HostListProvider, Protocol):
@@ -519,6 +522,10 @@ def _identify_connection(self, conn: Connection):
             cursor.execute(self._dialect.host_id_query)
             return cursor.fetchone()
 
+    def get_cluster_id(self):
+        self._initialize()
+        return self._cluster_id
+
     @dataclass()
     class ClusterIdSuggestion:
         cluster_id: str
@@ -646,3 +653,6 @@ def get_host_role(self, connection: Connection) -> HostRole:
     def identify_connection(self, connection: Optional[Connection]) -> Optional[HostInfo]:
         raise UnsupportedOperationError(
             Messages.get_formatted("ConnectionStringHostListProvider.UnsupportedMethod", "identify_connection"))
+
+    def get_cluster_id(self):
+        return "<none>"