feat: configure target driver dialect before connection (#61)

Author: karenc-bq

aws_wrapper/connection_provider.py

@@ -14,20 +14,22 @@
 
 from __future__ import annotations
 
+from logging import getLogger
 from threading import Lock
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Protocol
 
 if TYPE_CHECKING:
     from aws_wrapper.hostinfo import HostInfo, HostRole
     from aws_wrapper.pep249 import Connection
-    from aws_wrapper.utils.properties import Properties
-
-from typing import Callable, Dict, List, Optional, Protocol
+    from aws_wrapper.target_driver_dialect import TargetDriverDialect
 
 from aws_wrapper.errors import AwsWrapperError
 from aws_wrapper.hostselector import HostSelector, RandomHostSelector
 from aws_wrapper.plugin import CanReleaseResources
 from aws_wrapper.utils.messages import Messages
+from aws_wrapper.utils.properties import Properties, PropertiesUtils
+
+logger = getLogger(__name__)
 
 
 class ConnectionProvider(Protocol):
@@ -101,11 +103,11 @@ def release_resources():
 
 
 class DriverConnectionProvider(ConnectionProvider):
-
     _accepted_strategies: Dict[str, HostSelector] = {"random": RandomHostSelector()}
 
-    def __init__(self, connect_func: Callable):
+    def __init__(self, connect_func: Callable, target_driver_dialect: TargetDriverDialect):
         self._connect_func = connect_func
+        self._target_driver_dialect = target_driver_dialect
 
     def accepts_host_info(self, host_info: HostInfo, properties: Properties) -> bool:
         return True
@@ -122,12 +124,8 @@ def get_host_info_by_strategy(self, hosts: List[HostInfo], role: HostRole, strat
             return host_selector.get_host(hosts, role)
 
     def connect(self, host_info: HostInfo, properties: Properties) -> Connection:
-        # TODO: Behavior based on dialects
-        prop_copy = properties.copy()
-
-        prop_copy["host"] = host_info.host
-
-        if host_info.is_port_specified():
-            prop_copy["port"] = str(host_info.port)
+        prepared_properties = self._target_driver_dialect.prepare_connect_info(host_info, properties)
+        logger.debug(
+            f"Connecting to {host_info.host} with properties: {PropertiesUtils.log_properties(prepared_properties)}")
 
-        return self._connect_func(**prop_copy)
+        return self._connect_func(**prepared_properties)

aws_wrapper/dialect.py

@@ -27,6 +27,7 @@
                                           WrapperProperties)
 from aws_wrapper.utils.rdsutils import RdsUtils
 from .exceptions import ExceptionHandler, PgExceptionHandler
+from .target_driver_dialect import TargetDriverDialectCodes
 from .utils.cache_map import CacheMap
 from .utils.messages import Messages
 
@@ -55,7 +56,7 @@ def from_string(value: str) -> DialectCode:
             raise AwsWrapperError(Messages.get_formatted("DialectCode.InvalidStringValue", value))
 
 
-class DatabaseType(Enum):
+class TargetDriverType(Enum):
     MYSQL = auto()
     POSTGRES = auto()
     MARIADB = auto()
@@ -121,7 +122,7 @@ def exception_handler(self) -> Optional[ExceptionHandler]:
 
 
 class DialectProvider(Protocol):
-    def get_dialect(self, props: Properties) -> Optional[Dialect]:
+    def get_dialect(self, driver_dialect: str, props: Properties) -> Optional[Dialect]:
         """
         Returns the dialect identified by analyzing the AwsWrapperProperties.DIALECT property (if set) or the target
         driver method
@@ -470,7 +471,7 @@ def reset_custom_dialect(self):
     def reset_endpoint_cache(self):
         self._known_endpoint_dialects.clear()
 
-    def get_dialect(self, props: Properties) -> Optional[Dialect]:
+    def get_dialect(self, driver_dialect: str, props: Properties) -> Optional[Dialect]:
         self._can_update = False
         self._dialect = None
 
@@ -499,8 +500,8 @@ def get_dialect(self, props: Properties) -> Optional[Dialect]:
                 raise AwsWrapperError(Messages.get_formatted("Dialect.UnknownDialectCode", str(dialect_code)))
 
         host: str = props["host"]
-        database_type: DatabaseType = self._get_database_type()
-        if database_type is DatabaseType.MYSQL:
+        target_driver_type: TargetDriverType = self._get_target_driver_type(driver_dialect)
+        if target_driver_type is TargetDriverType.MYSQL:
             rds_type = self._rds_helper.identify_rds_type(host)
             if rds_type.is_rds_cluster:
                 self._dialect_code = DialectCode.AURORA_MYSQL
@@ -518,7 +519,7 @@ def get_dialect(self, props: Properties) -> Optional[Dialect]:
             self._log_current_dialect()
             return self._dialect
 
-        if database_type is DatabaseType.POSTGRES:
+        if target_driver_type is TargetDriverType.POSTGRES:
             rds_type = self._rds_helper.identify_rds_type(host)
             if rds_type.is_rds_cluster:
                 self._dialect_code = DialectCode.AURORA_PG
@@ -536,7 +537,14 @@ def get_dialect(self, props: Properties) -> Optional[Dialect]:
             self._log_current_dialect()
             return self._dialect
 
-        if database_type is DatabaseType.MARIADB:
+        if target_driver_type is TargetDriverType.MARIADB:
+            rds_type = self._rds_helper.identify_rds_type(host)
+            if rds_type.is_rds_cluster:
+                # Aurora MariaDB doesn't exist.
+                # If this is a cluster endpoint then user is trying to connect to AMS via the MariaDB driver.
+                self._dialect_code = DialectCode.AURORA_MYSQL
+                self._dialect = self._known_dialects_by_code.get(DialectCode.AURORA_MYSQL)
+                return self._dialect
             self._can_update = True
             self._dialect_code = DialectCode.MARIADB
             self._dialect = self._known_dialects_by_code.get(DialectCode.MARIADB)
@@ -549,9 +557,15 @@ def get_dialect(self, props: Properties) -> Optional[Dialect]:
         self._log_current_dialect()
         return self._dialect
 
-    def _get_database_type(self) -> DatabaseType:
-        # TODO: Add logic to identify database based on target driver connect info
-        return DatabaseType.POSTGRES
+    def _get_target_driver_type(self, driver_dialect: str) -> TargetDriverType:
+        if driver_dialect == TargetDriverDialectCodes.PSYCOPG:
+            return TargetDriverType.POSTGRES
+        if driver_dialect == TargetDriverDialectCodes.MYSQL_CONNECTOR_PYTHON:
+            return TargetDriverType.MYSQL
+        if driver_dialect == TargetDriverDialectCodes.MARIADB_CONNECTOR_PYTHON:
+            return TargetDriverType.MARIADB
+
+        return TargetDriverType.CUSTOM
 
     def query_for_dialect(self, url: str, host_info: Optional[HostInfo], conn: Connection) -> Optional[Dialect]:
         if not self._can_update:

aws_wrapper/failover_plugin.py

@@ -41,16 +41,14 @@
 from aws_wrapper.utils.properties import Properties, WrapperProperties
 from aws_wrapper.utils.rds_url_type import RdsUrlType
 from aws_wrapper.utils.rdsutils import RdsUtils
-from aws_wrapper.utils.utils import SubscribedMethodUtils
 from aws_wrapper.writer_failover_handler import (WriterFailoverHandler,
                                                  WriterFailoverHandlerImpl)
 
 logger = getLogger(__name__)
 
 
 class FailoverPlugin(Plugin):
-    _SUBSCRIBED_METHODS: Set[str] = {*SubscribedMethodUtils.NETWORK_BOUND_METHODS,
-                                     "init_host_provider",
+    _SUBSCRIBED_METHODS: Set[str] = {"init_host_provider",
                                      "connect",
                                      "force_connect",
                                      "notify_host_list_changed"}
@@ -69,6 +67,7 @@ def __init__(self, plugin_service: PluginService, props: Properties):
         self._last_exception: Optional[Exception] = None
         self._rds_utils = RdsUtils()
         self._rds_url_type: RdsUrlType = self._rds_utils.identify_rds_type(self._properties.get("host"))
+        FailoverPlugin._SUBSCRIBED_METHODS.add(*self._plugin_service.network_bound_methods)
 
     def init_host_provider(
             self,

aws_wrapper/generic_target_driver_dialect.py

@@ -0,0 +1,62 @@
+#  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 abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Callable, Set
+
+from aws_wrapper.target_driver_dialect_codes import TargetDriverDialectCodes
+from aws_wrapper.utils.properties import Properties, PropertiesUtils
+
+if TYPE_CHECKING:
+    from aws_wrapper.hostinfo import HostInfo
+
+
+class TargetDriverDialect(ABC):
+    _dialect_code: str = TargetDriverDialectCodes.GENERIC
+    _network_bound_methods: Set[str] = {"*"}
+
+    @property
+    def dialect_code(self) -> str:
+        return self._dialect_code
+
+    @property
+    def network_bound_methods(self) -> Set[str]:
+        return self._network_bound_methods
+
+    @abstractmethod
+    def is_dialect(self, conn: Callable) -> bool:
+        pass
+
+    @abstractmethod
+    def prepare_connect_info(self, host_info: HostInfo, props: Properties) -> Properties:
+        pass
+
+
+class GenericTargetDriverDialect(TargetDriverDialect):
+
+    def is_dialect(self, conn: Callable) -> bool:
+        return True
+
+    def prepare_connect_info(self, host_info: HostInfo, props: Properties) -> Properties:
+        prop_copy: Properties = Properties(props.copy())
+
+        prop_copy["host"] = host_info.host
+
+        if host_info.is_port_specified():
+            prop_copy["port"] = str(host_info.port)
+
+        PropertiesUtils.remove_wrapper_props(prop_copy)
+        return props

aws_wrapper/host_monitoring_plugin.py

@@ -40,7 +40,7 @@
 from aws_wrapper.utils.properties import Properties, WrapperProperties
 from aws_wrapper.utils.rdsutils import RdsUtils
 from aws_wrapper.utils.timeout import timeout
-from aws_wrapper.utils.utils import QueueUtils, SubscribedMethodUtils
+from aws_wrapper.utils.utils import QueueUtils
 
 logger = getLogger(__name__)
 
@@ -92,7 +92,7 @@ def execute(self, target: object, method_name: str, execute_func: Callable, *arg
             raise AwsWrapperError(Messages.get_formatted("HostMonitoringPlugin.NullHostInfoForMethod", method_name))
 
         is_enabled = WrapperProperties.FAILURE_DETECTION_ENABLED.get_bool(self._props)
-        if not is_enabled or method_name not in SubscribedMethodUtils.NETWORK_BOUND_METHODS:
+        if not is_enabled or not self._plugin_service.is_network_bound_method(method_name):
             return execute_func()
 
         failure_detection_time_ms = WrapperProperties.FAILURE_DETECTION_TIME_MS.get_int(self._props)

aws_wrapper/mariadb_target_driver_dialect.py

@@ -0,0 +1,42 @@
+#  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 inspect import signature
+from typing import Callable, Set
+
+from aws_wrapper.generic_target_driver_dialect import \
+    GenericTargetDriverDialect
+from aws_wrapper.target_driver_dialect_codes import TargetDriverDialectCodes
+
+
+class MariaDBTargetDriverDialect(GenericTargetDriverDialect):
+    TARGET_DRIVER = "Mariadb"
+
+    _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()",
+    }
+
+    def is_dialect(self, conn: Callable) -> bool:
+        return MariaDBTargetDriverDialect.TARGET_DRIVER in str(signature(conn))

aws_wrapper/mysql_target_driver_dialect.py

@@ -0,0 +1,41 @@
+#  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 inspect import signature
+from typing import Callable, Set
+
+from aws_wrapper.generic_target_driver_dialect import \
+    GenericTargetDriverDialect
+from aws_wrapper.target_driver_dialect_codes import TargetDriverDialectCodes
+
+
+class MySQLTargetDriverDialect(GenericTargetDriverDialect):
+    TARGET_DRIVER = "MySQL"
+
+    _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()",
+    }
+
+    def is_dialect(self, conn: Callable) -> bool:
+        return MySQLTargetDriverDialect.TARGET_DRIVER in str(signature(conn))

aws_wrapper/pg_target_driver_dialect.py

@@ -0,0 +1,40 @@
+#  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 inspect import signature
+from typing import Callable, Set
+
+from aws_wrapper.target_driver_dialect import GenericTargetDriverDialect
+from aws_wrapper.target_driver_dialect_codes import TargetDriverDialectCodes
+
+
+class PgTargetDriverDialect(GenericTargetDriverDialect):
+    TARGET_DRIVER = "psycopg"
+
+    _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()",
+    }
+
+    def is_dialect(self, conn: Callable) -> bool:
+        return PgTargetDriverDialect.TARGET_DRIVER in str(signature(conn))