fix: DriverDialect.is_dialect may return false for different driver c…

…lasses (#523)
aws · May 8, 2024 · c51572f · c51572f
1 parent 66e3203
commit c51572f
Show file tree

Hide file tree

Showing 9 changed files with 16 additions and 12 deletions.
diff --git a/aws_advanced_python_wrapper/connection_provider.py b/aws_advanced_python_wrapper/connection_provider.py
@@ -25,9 +25,9 @@
 from threading import Lock
 
 from aws_advanced_python_wrapper.errors import AwsWrapperError
-from aws_advanced_python_wrapper.hostselector import (HostSelector,
-                                                      RandomHostSelector,
-                                                      RoundRobinHostSelector)
+from aws_advanced_python_wrapper.host_selector import (HostSelector,
+                                                       RandomHostSelector,
+                                                       RoundRobinHostSelector)
 from aws_advanced_python_wrapper.plugin import CanReleaseResources
 from aws_advanced_python_wrapper.utils.log import Logger
 from aws_advanced_python_wrapper.utils.messages import Messages

diff --git a/aws_advanced_python_wrapper/fastest_response_strategy_plugin.py b/aws_advanced_python_wrapper/fastest_response_strategy_plugin.py
@@ -23,7 +23,7 @@
                     Set, Tuple)
 
 from aws_advanced_python_wrapper.errors import AwsWrapperError
-from aws_advanced_python_wrapper.hostselector import RandomHostSelector
+from aws_advanced_python_wrapper.host_selector import RandomHostSelector
 from aws_advanced_python_wrapper.plugin import Plugin
 from aws_advanced_python_wrapper.utils.cache_map import CacheMap
 from aws_advanced_python_wrapper.utils.log import Logger

diff --git a/aws_advanced_python_wrapper/hostselector.py → aws_advanced_python_wrapper/host_selector.py b/aws_advanced_python_wrapper/hostselector.py → aws_advanced_python_wrapper/host_selector.py
diff --git a/aws_advanced_python_wrapper/mysql_driver_dialect.py b/aws_advanced_python_wrapper/mysql_driver_dialect.py
@@ -63,7 +63,9 @@ class MySQLDriverDialect(DriverDialect):
     }
 
     def is_dialect(self, connect_func: Callable) -> bool:
-        return MySQLDriverDialect.TARGET_DRIVER_CODE in str(signature(connect_func))
+        if MySQLDriverDialect.TARGET_DRIVER_CODE not in str(signature(connect_func)):
+            return MySQLDriverDialect.TARGET_DRIVER_CODE.lower() in (connect_func.__module__ + connect_func.__qualname__).lower()
+        return True
 
     def is_closed(self, conn: Connection) -> bool:
         if isinstance(conn, CMySQLConnection) or isinstance(conn, MySQLConnection):

diff --git a/aws_advanced_python_wrapper/pg_driver_dialect.py b/aws_advanced_python_wrapper/pg_driver_dialect.py
@@ -58,7 +58,9 @@ class PgDriverDialect(DriverDialect):
     }
 
     def is_dialect(self, connect_func: Callable) -> bool:
-        return PgDriverDialect.TARGET_DRIVER_CODE in str(signature(connect_func))
+        if PgDriverDialect.TARGET_DRIVER_CODE not in str(signature(connect_func)):
+            return PgDriverDialect.TARGET_DRIVER_CODE.lower() in (connect_func.__module__ + connect_func.__qualname__).lower()
+        return True
 
     def is_closed(self, conn: Connection) -> bool:
         if isinstance(conn, psycopg.Connection):

diff --git a/aws_advanced_python_wrapper/sql_alchemy_connection_provider.py b/aws_advanced_python_wrapper/sql_alchemy_connection_provider.py
@@ -25,9 +25,9 @@
 
 from aws_advanced_python_wrapper.connection_provider import ConnectionProvider
 from aws_advanced_python_wrapper.errors import AwsWrapperError
-from aws_advanced_python_wrapper.hostselector import (HostSelector,
-                                                      RandomHostSelector,
-                                                      RoundRobinHostSelector)
+from aws_advanced_python_wrapper.host_selector import (HostSelector,
+                                                       RandomHostSelector,
+                                                       RoundRobinHostSelector)
 from aws_advanced_python_wrapper.plugin import CanReleaseResources
 from aws_advanced_python_wrapper.utils.messages import Messages
 from aws_advanced_python_wrapper.utils.properties import (Properties,

diff --git a/tests/unit/test_host_selector.py b/tests/unit/test_host_selector.py
@@ -16,8 +16,8 @@
 import pytest
 
 from aws_advanced_python_wrapper import pep249
+from aws_advanced_python_wrapper.host_selector import RandomHostSelector
 from aws_advanced_python_wrapper.hostinfo import HostInfo, HostRole
-from aws_advanced_python_wrapper.hostselector import RandomHostSelector
 
 
 def test_random_host_selection_strategy():

diff --git a/tests/unit/test_random_host_selector.py b/tests/unit/test_random_host_selector.py
@@ -15,8 +15,8 @@
 import pytest
 
 from aws_advanced_python_wrapper.host_availability import HostAvailability
+from aws_advanced_python_wrapper.host_selector import RandomHostSelector
 from aws_advanced_python_wrapper.hostinfo import HostInfo, HostRole
-from aws_advanced_python_wrapper.hostselector import RandomHostSelector
 from aws_advanced_python_wrapper.utils.properties import Properties
 
 HOST_ROLE = HostRole.READER

diff --git a/tests/unit/test_round_robin_host_selector.py b/tests/unit/test_round_robin_host_selector.py
@@ -16,8 +16,8 @@
 
 from aws_advanced_python_wrapper.errors import AwsWrapperError
 from aws_advanced_python_wrapper.host_availability import HostAvailability
+from aws_advanced_python_wrapper.host_selector import RoundRobinHostSelector
 from aws_advanced_python_wrapper.hostinfo import HostInfo, HostRole
-from aws_advanced_python_wrapper.hostselector import RoundRobinHostSelector
 from aws_advanced_python_wrapper.utils.properties import (Properties,
                                                           WrapperProperties)