feat: efm v2 support

JuanLeee · JuanLeee · commit 55dae2ff65b6 · 2025-08-08T10:43:32.000-07:00
# Conflicts:
#	aws_advanced_python_wrapper/plugin_service.py
#	aws_advanced_python_wrapper/utils/atomic.py
diff --git a/aws_advanced_python_wrapper/host_monitoring_v2_plugin.py b/aws_advanced_python_wrapper/host_monitoring_v2_plugin.py
diff --git a/aws_advanced_python_wrapper/plugin_service.py b/aws_advanced_python_wrapper/plugin_service.py
@@ -75,6 +75,8 @@
     HostListProviderService, StaticHostListProvider)
 from aws_advanced_python_wrapper.host_monitoring_plugin import \
     HostMonitoringPluginFactory
+from aws_advanced_python_wrapper.host_monitoring_v2_plugin import \
+    HostMonitoringV2PluginFactory
 from aws_advanced_python_wrapper.hostinfo import HostInfo, HostRole
 from aws_advanced_python_wrapper.iam_plugin import IamAuthPluginFactory
 from aws_advanced_python_wrapper.plugin import CanReleaseResources
@@ -755,6 +757,7 @@ class PluginManager(CanReleaseResources):
         "aws_secrets_manager": AwsSecretsManagerPluginFactory,
         "aurora_connection_tracker": AuroraConnectionTrackerPluginFactory,
         "host_monitoring": HostMonitoringPluginFactory,
+        "host_monitoring_v2": HostMonitoringV2PluginFactory,
         "failover": FailoverPluginFactory,
         "read_write_splitting": ReadWriteSplittingPluginFactory,
         "fastest_response_strategy": FastestResponseStrategyPluginFactory,
@@ -783,6 +786,7 @@ class PluginManager(CanReleaseResources):
         ReadWriteSplittingPluginFactory: 300,
         FailoverPluginFactory: 400,
         HostMonitoringPluginFactory: 500,
+        HostMonitoringV2PluginFactory: 510,
         BlueGreenPluginFactory: 550,
         FastestResponseStrategyPluginFactory: 600,
         IamAuthPluginFactory: 700,
diff --git a/aws_advanced_python_wrapper/resources/aws_advanced_python_wrapper_messages.properties b/aws_advanced_python_wrapper/resources/aws_advanced_python_wrapper_messages.properties
@@ -176,6 +176,39 @@ HostMonitoringPlugin.ConfigurationNotSupported=[HostMonitoringPlugin] Aborting c
 HostMonitoringPlugin.UnableToIdentifyConnection=[HostMonitoringPlugin] Unable to identify the connected database instance: '{}', please ensure the correct host list provider is specified. The host list provider in use is: '{}'.
 HostMonitoringPlugin.UnavailableHost=[HostMonitoringPlugin] Host '{}' is unavailable.
 
+HostMonitoringV2Plugin.ActivatedMonitoring=[HostMonitoringV2Plugin] Executing method '{}', monitoring is activated.
+HostMonitoringV2Plugin.ClusterEndpointHostInfo=[HostMonitoringV2Plugin] The HostInfo to monitor is associated with a cluster endpoint. The plugin will attempt to identify the connected database instance.
+HostMonitoringV2Plugin.ErrorIdentifyingConnection=[HostMonitoringV2Plugin] An error occurred while identifying the connection database instance: '{}'.
+HostMonitoringV2Plugin.MonitoringDeactivated=[HostMonitoringV2Plugin] Monitoring deactivated for method '{}'.
+HostMonitoringV2Plugin.ConnectionNone=[HostMonitoringV2Plugin] Attempted to execute method '{}' but the current connection is None.
+HostMonitoringV2Plugin.HostInfoNone=[HostMonitoringV2Plugin] Could not find HostInfo to monitor for the current connection.
+HostMonitoringV2Plugin.HostInfoNoneForMethod=[HostMonitoringV2Plugin] Attempted to execute method '{}' but the current host info is None.
+HostMonitoringV2Plugin.ConfigurationNotSupported=[HostMonitoringV2Plugin] Aborting connections from a separate thread is not supported for the detected driver dialect: '{}'. The EFM plugin requires this feature to be supported.
+HostMonitoringV2Plugin.UnableToIdentifyConnection=[HostMonitoringV2Plugin] Unable to identify the connected database instance: '{}', please ensure the correct host list provider is specified. The host list provider in use is: '{}'.
+HostMonitoringV2Plugin.UnavailableHost=[HostMonitoringV2Plugin] Host '{}' is unavailable.
+
+HostMonitorV2.ContextNullWarning=[HostMonitorV2] Parameter 'context' should not be null.
+HostMonitorV2.InterruptedExceptionDuringMonitoring=[HostMonitorV2] Monitoring thread for node {} was interrupted.
+HostMonitorV2.ExceptionDuringMonitoringContinue=[HostMonitorV2] Continuing monitoring after unhandled exception was thrown in monitoring thread for node '{}'.
+HostMonitorV2.ExceptionDuringMonitoringStop=[HostMonitorV2] Stopping monitoring after unhandled exception was thrown in monitoring thread for node '{}'.
+HostMonitorV2.MonitorIsStopped=[HostMonitorV2] Monitoring was already stopped for node '{}'.
+HostMonitorV2.Stopped=[HostMonitorV2] Stopped monitoring thread for node '{}'.
+HostMonitorV2.StartMonitoringThreadNewContext=[HostMonitorV2] Start monitoring thread for checking new contexts for '{}'.
+HostMonitorV2.StopMonitoringThreadNewContext=[HostMonitorV2] Stop monitoring thread for checking new contexts for '{}'.
+HostMonitorV2.StartMonitoringThread=[HostMonitorV2] Start monitoring thread for '{}'.
+HostMonitorV2.StopMonitoringThread=[HostMonitorV2] Stop monitoring thread for '{}'.
+HostMonitorV2.OpeningMonitoringConnection=[HostMonitorV2] Opening a monitoring connection to '{}'
+HostMonitorV2.OpenedMonitoringConnection=[HostMonitorV2] Opened monitoring connection: '{}'
+HostMonitorV2.ExceptionAbortingConnection=[HostMonitorV2] Exception during aborting connection: '{}'
+HostMonitorV2.HostDead=[HostMonitorV2] Host '{}' is *dead*.
+HostMonitorV2.HostNotResponding=[HostMonitorV2] Host '{}' is not *responding* '{}'.
+HostMonitorV2.HostAlive=[HostMonitorV2] Host '{}' is *alive*.
+
+MonitorServiceV2.ExceptionAbortingConnection=[MonitorServiceV2] Exception during aborting connection: '{}'
+MonitorServiceV2.HostDead=[MonitorServiceV2] Host '{}' is *dead*.
+MonitorServiceV2.HostNotResponding=[MonitorServiceV2] Host '{}' is not *responding* '{}'.
+MonitorServiceV2.HostAlive=[MonitorServiceV2] Host '{}' is *alive*.
+
 HostSelector.NoEligibleHost=[HostSelector] No Eligible Hosts Found.
 HostSelector.NoHostsMatchingRole=[HostSelector] No hosts were found matching the requested role: '{}'.
 
diff --git a/aws_advanced_python_wrapper/utils/atomic.py b/aws_advanced_python_wrapper/utils/atomic.py
@@ -13,6 +13,7 @@
 #  limitations under the License.
 
 from threading import Lock
+from typing import Any
 
 
 class AtomicInt:
@@ -59,3 +60,31 @@ def compare_and_set(self, expected_value: int, new_value: int) -> bool:
                 self._value = new_value
                 return True
             return False
+
+
+class AtomicBoolean:
+    def __init__(self, initial_value: bool):
+        self._value = initial_value
+        self._lock: Lock = Lock()
+
+    def get(self) -> bool:
+        with self._lock:
+            return self._value
+
+    def set(self, value: bool) -> None:
+        with self._lock:
+            self._value = value
+
+
+class AtomicReference:
+    def __init__(self, initial_value):
+        self._value = initial_value
+        self._lock: Lock = Lock()
+
+    def get(self) -> Any:
+        with self._lock:
+            return self._value
+
+    def set(self, new_value: Any) -> None:
+        with self._lock:
+            self._value = new_value
diff --git a/tests/integration/container/test_aurora_failover.py b/tests/integration/container/test_aurora_failover.py
@@ -116,6 +116,7 @@ def test_fail_from_writer_to_new_writer_fail_on_connection_bound_object_invocati
             assert aurora_utility.is_db_instance_writer(current_connection_id) is True
             assert current_connection_id != initial_writer_id
 
+    @pytest.mark.parametrize("plugins", ["failover,host_monitoring", "failover,host_monitoring_v2"])
     @enable_on_features([TestEnvironmentFeatures.NETWORK_OUTAGES_ENABLED,
                          TestEnvironmentFeatures.ABORT_CONNECTION_SUPPORTED])
     def test_fail_from_reader_to_writer(
@@ -124,12 +125,13 @@ def test_fail_from_reader_to_writer(
             test_driver: TestDriver,
             conn_utils,
             proxied_props,
-            aurora_utility):
+            aurora_utility,
+            plugins):
         target_driver_connect = DriverHelper.get_connect_func(test_driver)
         reader: TestInstanceInfo = test_environment.get_proxy_instances()[1]
         writer_id: str = test_environment.get_proxy_writer().get_instance_id()
 
-        proxied_props["plugins"] = "failover,host_monitoring"
+        proxied_props["plugins"] = plugins
         with AwsWrapperConnection.connect(
                 target_driver_connect,
                 **conn_utils.get_proxy_connect_params(reader.get_host()),
diff --git a/tests/integration/container/test_basic_connectivity.py b/tests/integration/container/test_basic_connectivity.py
@@ -126,15 +126,32 @@ def test_proxied_wrapper_connection_failed(
             # That is expected exception. Test pass.
             assert True
 
+    @pytest.mark.parametrize("plugins", ["host_monitoring", "host_monitoring_v2"])
     @enable_on_num_instances(min_instances=2)
     @enable_on_deployments([DatabaseEngineDeployment.AURORA, DatabaseEngineDeployment.RDS_MULTI_AZ_CLUSTER])
     @enable_on_features([TestEnvironmentFeatures.ABORT_CONNECTION_SUPPORTED])
-    def test_wrapper_connection_reader_cluster_with_efm_enabled(self, test_driver: TestDriver, conn_utils):
+    def test_wrapper_connection_reader_cluster_with_efm_enabled(self, test_driver: TestDriver, conn_utils, plugins):
         target_driver_connect = DriverHelper.get_connect_func(test_driver)
         conn = AwsWrapperConnection.connect(
             target_driver_connect,
             **conn_utils.get_connect_params(conn_utils.reader_cluster_host),
-            plugins="host_monitoring", connect_timeout=10)
+            plugins=plugins, connect_timeout=10)
+        cursor = conn.cursor()
+        cursor.execute("SELECT 1")
+        result = cursor.fetchone()
+        assert 1 == result[0]
+
+        conn.close()
+
+    @enable_on_num_instances(min_instances=2)
+    @enable_on_deployments([DatabaseEngineDeployment.AURORA, DatabaseEngineDeployment.RDS_MULTI_AZ_CLUSTER])
+    @enable_on_features([TestEnvironmentFeatures.ABORT_CONNECTION_SUPPORTED])
+    def test_wrapper_connection_reader_cluster_with_efmv2_enabled(self, test_driver: TestDriver, conn_utils):
+        target_driver_connect = DriverHelper.get_connect_func(test_driver)
+        conn = AwsWrapperConnection.connect(
+            target_driver_connect,
+            **conn_utils.get_connect_params(conn_utils.reader_cluster_host),
+            plugins="host_monitoring_v2", connect_timeout=10)
         cursor = conn.cursor()
         cursor.execute("SELECT 1")
         result = cursor.fetchone()
diff --git a/tests/integration/container/test_failover_performance.py b/tests/integration/container/test_failover_performance.py
@@ -131,8 +131,9 @@ def props(self):
 
         return props
 
+    @pytest.mark.parametrize("plugins", ["host_monitoring", "host_monitoring_v2"])
     def test_failure_detection_time_efm(self, test_environment: TestEnvironment, test_driver: TestDriver, conn_utils,
-                                        props: Properties):
+                                        props: Properties, plugins):
         enhanced_failure_monitoring_perf_data_list: List[PerfStatBase] = []
         target_driver_connect_func = DriverHelper.get_connect_func(test_driver)
         try:
@@ -147,7 +148,7 @@ def test_failure_detection_time_efm(self, test_environment: TestEnvironment, tes
                 WrapperProperties.FAILURE_DETECTION_TIME_MS.set(props, str(detection_time))
                 WrapperProperties.FAILURE_DETECTION_INTERVAL_MS.set(props, str(detection_interval))
                 WrapperProperties.FAILURE_DETECTION_COUNT.set(props, str(detection_count))
-                WrapperProperties.PLUGINS.set(props, "host_monitoring")
+                WrapperProperties.PLUGINS.set(props, plugins)
 
                 data: PerfStatMonitoring = PerfStatMonitoring()
                 self._measure_performance(test_environment, target_driver_connect_func, conn_utils, sleep_delay_sec, props, data)
@@ -162,8 +163,9 @@ def test_failure_detection_time_efm(self, test_environment: TestEnvironment, tes
                 f"FailureDetectionPerformanceResults_EnhancedMonitoringEnabled.csv",
                 TestPerformance.PERF_STAT_MONITORING_HEADER, enhanced_failure_monitoring_perf_data_list)
 
+    @pytest.mark.parametrize("plugins", ["failover,host_monitoring", "failover,host_monitoring_v2"])
     def test_failure_detection_time_failover_and_efm(self, test_environment: TestEnvironment, test_driver: TestDriver, conn_utils,
-                                                     props: Properties):
+                                                     props: Properties, plugins):
         enhanced_failure_monitoring_perf_data_list: List[PerfStatBase] = []
         try:
             for i in range(len(TestPerformance.failure_detection_time_params)):
@@ -177,7 +179,7 @@ def test_failure_detection_time_failover_and_efm(self, test_environment: TestEnv
                 WrapperProperties.FAILURE_DETECTION_TIME_MS.set(props, str(detection_time))
                 WrapperProperties.FAILURE_DETECTION_INTERVAL_MS.set(props, str(detection_interval))
                 WrapperProperties.FAILURE_DETECTION_COUNT.set(props, str(detection_count))
-                WrapperProperties.PLUGINS.set(props, "failover,host_monitoring")
+                WrapperProperties.PLUGINS.set(props, plugins)
                 WrapperProperties.FAILOVER_TIMEOUT_SEC.set(props, TestPerformance.PERF_FAILOVER_TIMEOUT_SEC)
                 WrapperProperties.FAILOVER_MODE.set(props, "strict_reader")
 
diff --git a/tests/integration/container/test_read_write_splitting.py b/tests/integration/container/test_read_write_splitting.py
@@ -346,14 +346,16 @@ def test_failover_to_new_writer__switch_read_only(
         current_id = rds_utils.query_instance_id(conn)
         assert new_writer_id == current_id
 
+    @pytest.mark.parametrize("plugins", ["read_write_splitting,failover,host_monitoring",
+                                         "read_write_splitting,failover,host_monitoring_v2"])
     @enable_on_features([TestEnvironmentFeatures.NETWORK_OUTAGES_ENABLED,
                          TestEnvironmentFeatures.ABORT_CONNECTION_SUPPORTED])
     @enable_on_num_instances(min_instances=3)
     @disable_on_engines([DatabaseEngine.MYSQL])
     def test_failover_to_new_reader__switch_read_only(
             self, test_environment: TestEnvironment, test_driver: TestDriver,
-            proxied_failover_props, conn_utils, rds_utils):
-        WrapperProperties.PLUGINS.set(proxied_failover_props, "read_write_splitting,failover,host_monitoring")
+            proxied_failover_props, conn_utils, rds_utils, plugins):
+        WrapperProperties.PLUGINS.set(proxied_failover_props, plugins)
         WrapperProperties.FAILOVER_MODE.set(proxied_failover_props, "reader-or-writer")
 
         target_driver_connect = DriverHelper.get_connect_func(test_driver)
@@ -394,14 +396,16 @@ def test_failover_to_new_reader__switch_read_only(
         current_id = rds_utils.query_instance_id(conn)
         assert other_reader_id == current_id
 
+    @pytest.mark.parametrize("plugins", ["read_write_splitting,failover,host_monitoring",
+                                         "read_write_splitting,failover,host_monitoring_v2"])
     @enable_on_features([TestEnvironmentFeatures.NETWORK_OUTAGES_ENABLED,
                          TestEnvironmentFeatures.ABORT_CONNECTION_SUPPORTED])
     @enable_on_num_instances(min_instances=3)
     @disable_on_engines([DatabaseEngine.MYSQL])
     def test_failover_reader_to_writer__switch_read_only(
             self, test_environment: TestEnvironment, test_driver: TestDriver,
-            proxied_failover_props, conn_utils, rds_utils):
-        WrapperProperties.PLUGINS.set(proxied_failover_props, "read_write_splitting,failover,host_monitoring")
+            proxied_failover_props, conn_utils, rds_utils, plugins):
+        WrapperProperties.PLUGINS.set(proxied_failover_props, plugins)
         target_driver_connect = DriverHelper.get_connect_func(test_driver)
         conn = AwsWrapperConnection.connect(
             target_driver_connect, **conn_utils.get_proxy_connect_params(), **proxied_failover_props)
@@ -513,17 +517,19 @@ def test_pooled_connection__cluster_url_failover(
         new_driver_conn = conn.target_connection
         assert initial_driver_conn is not new_driver_conn
 
+    @pytest.mark.parametrize("plugins", ["read_write_splitting,failover,host_monitoring",
+                                         "read_write_splitting,failover,host_monitoring_v2"])
     @enable_on_features([TestEnvironmentFeatures.FAILOVER_SUPPORTED, TestEnvironmentFeatures.NETWORK_OUTAGES_ENABLED,
                          TestEnvironmentFeatures.ABORT_CONNECTION_SUPPORTED])
     @disable_on_engines([DatabaseEngine.MYSQL])
     def test_pooled_connection__failover_failed(
             self, test_environment: TestEnvironment, test_driver: TestDriver,
-            rds_utils, conn_utils, proxied_failover_props):
+            rds_utils, conn_utils, proxied_failover_props, plugins):
         writer_host = test_environment.get_writer().get_host()
         provider = SqlAlchemyPooledConnectionProvider(lambda _, __: {"pool_size": 1}, None, lambda host_info, props: writer_host in host_info.host)
         ConnectionProviderManager.set_connection_provider(provider)
 
-        WrapperProperties.PLUGINS.set(proxied_failover_props, "read_write_splitting,failover,host_monitoring")
+        WrapperProperties.PLUGINS.set(proxied_failover_props, plugins)
         WrapperProperties.FAILOVER_TIMEOUT_SEC.set(proxied_failover_props, "1")
         WrapperProperties.FAILURE_DETECTION_TIME_MS.set(proxied_failover_props, "1000")
         WrapperProperties.FAILURE_DETECTION_COUNT.set(proxied_failover_props, "1")
diff --git a/tests/unit/test_host_monitor_v2_plugin.py b/tests/unit/test_host_monitor_v2_plugin.py
