fix: subscribe methods shouldn't end with () and topology timeout (#70)

Author: karenc-bq

aws_wrapper/errors.py

@@ -21,3 +21,7 @@ class AwsWrapperError(Error):
 
 class FailoverError(AwsWrapperError):
     __module__ = "aws_wrapper"
+
+
+class QueryTimeoutError(AwsWrapperError):
+    __module__ = "aws_wrapper"

aws_wrapper/exceptions.py

@@ -16,6 +16,8 @@
 
 from typing import TYPE_CHECKING
 
+from aws_wrapper.errors import QueryTimeoutError
+
 if TYPE_CHECKING:
     from aws_wrapper.dialect import Dialect
 
@@ -55,7 +57,8 @@ class PgExceptionHandler(ExceptionHandler):
     _PAM_AUTHENTICATION_FAILED_MSG = "PAM authentication failed"
 
     def is_network_exception(self, error: Optional[Exception] = None, sql_state: Optional[str] = None) -> bool:
-        # TODO: verify implementation
+        if isinstance(error, QueryTimeoutError):
+            return True
 
         if sql_state:
             if sql_state in self._NETWORK_ERRORS:

aws_wrapper/generic_target_driver_dialect.py

@@ -59,4 +59,4 @@ def prepare_connect_info(self, host_info: HostInfo, props: Properties) -> Proper
             prop_copy["port"] = str(host_info.port)
 
         PropertiesUtils.remove_wrapper_props(prop_copy)
-        return props
+        return prop_copy

aws_wrapper/host_list_provider.py

@@ -28,7 +28,7 @@
     from aws_wrapper.plugin_service import PluginService
 
 from aws_wrapper.dialect import Dialect, TopologyAwareDatabaseDialect
-from aws_wrapper.errors import AwsWrapperError
+from aws_wrapper.errors import AwsWrapperError, QueryTimeoutError
 from aws_wrapper.hostinfo import HostAvailability, HostInfo, HostRole
 from aws_wrapper.pep249 import Connection, Cursor, Error, ProgrammingError
 from aws_wrapper.plugin import Plugin, PluginFactory
@@ -238,7 +238,7 @@ def _get_topology(self, conn: Optional[Connection], force_update: bool = False)
                         self._suggest_cluster_id(hosts)
                     return AuroraHostListProvider.FetchTopologyResult(hosts, False)
             except Exception as e:
-                raise AwsWrapperError(Messages.get("AuroraHostListProvider.TopologyTimeout")) from e
+                raise QueryTimeoutError(Messages.get("AuroraHostListProvider.TopologyTimeout")) from e
 
         if cached_hosts:
             return AuroraHostListProvider.FetchTopologyResult(cached_hosts, True)

aws_wrapper/mariadb_target_driver_dialect.py

@@ -27,15 +27,15 @@ class MariaDBTargetDriverDialect(GenericTargetDriverDialect):
 
     _dialect_code: str = TargetDriverDialectCodes.MARIADB_CONNECTOR_PYTHON
     _network_bound_methods: Set[str] = {
-        "Connection.commit()",
-        "Connection.rollback()",
-        "Connection.cursor()",
-        "Cursor.callproc()",
-        "Cursor.execute()",
-        "Cursor.fetchone()",
-        "Cursor.fetchmany()",
-        "Cursor.fetchall()",
-        "Cursor.nextset()",
+        "Connection.commit",
+        "Connection.rollback",
+        "Connection.cursor",
+        "Cursor.callproc",
+        "Cursor.execute",
+        "Cursor.fetchone",
+        "Cursor.fetchmany",
+        "Cursor.fetchall",
+        "Cursor.nextset",
     }
 
     def is_dialect(self, conn: Callable) -> bool:

aws_wrapper/mysql_target_driver_dialect.py

@@ -27,14 +27,14 @@ class MySQLTargetDriverDialect(GenericTargetDriverDialect):
 
     _dialect_code: str = TargetDriverDialectCodes.MYSQL_CONNECTOR_PYTHON
     _network_bound_methods: Set[str] = {
-        "Connection.commit()",
-        "Connection.rollback()",
-        "Connection.cursor()",
-        "Cursor.close()",
-        "Cursor.execute()",
-        "Cursor.fetchone()",
-        "Cursor.fetchmany()",
-        "Cursor.fetchall()",
+        "Connection.commit",
+        "Connection.rollback",
+        "Connection.cursor",
+        "Cursor.close",
+        "Cursor.execute",
+        "Cursor.fetchone",
+        "Cursor.fetchmany",
+        "Cursor.fetchall",
     }
 
     def is_dialect(self, conn: Callable) -> bool:

aws_wrapper/pg_target_driver_dialect.py

@@ -26,14 +26,14 @@ class PgTargetDriverDialect(GenericTargetDriverDialect):
 
     _dialect_code: str = TargetDriverDialectCodes.PSYCOPG
     _network_bound_methods: Set[str] = {
-        "Connection.commit()",
-        "Connection.rollback()",
-        "Connection.cursor()",
-        "Cursor.callproc()",
-        "Cursor.execute()",
-        "Cursor.fetchone()",
-        "Cursor.fetchmany()",
-        "Cursor.fetchall()",
+        "Connection.commit",
+        "Connection.rollback",
+        "Connection.cursor",
+        "Cursor.callproc",
+        "Cursor.execute",
+        "Cursor.fetchone",
+        "Cursor.fetchmany",
+        "Cursor.fetchall",
     }
 
     def is_dialect(self, conn: Callable) -> bool: