add sql alchemy tortoise connection provider

Author: faviansamatha

aws_advanced_python_wrapper/driver_dialect.py

@@ -19,6 +19,7 @@
 if TYPE_CHECKING:
     from aws_advanced_python_wrapper.hostinfo import HostInfo
     from aws_advanced_python_wrapper.pep249 import Connection, Cursor
+    from types import ModuleType
 
 from abc import ABC
 from concurrent.futures import Executor, ThreadPoolExecutor, TimeoutError
@@ -164,3 +165,7 @@ def ping(self, conn: Connection) -> bool:
                 return True
         except Exception:
             return False
+    
+    def get_driver_module(self) -> ModuleType:
+        raise UnsupportedOperationError(
+            Messages.get_formatted("DriverDialect.UnsupportedOperationError", self._driver_name, "get_driver_module"))

aws_advanced_python_wrapper/driver_dialect_manager.py

@@ -24,8 +24,7 @@
 from aws_advanced_python_wrapper.errors import AwsWrapperError
 from aws_advanced_python_wrapper.utils.log import Logger
 from aws_advanced_python_wrapper.utils.messages import Messages
-from aws_advanced_python_wrapper.utils.properties import (Properties,
-                                                          WrapperProperties)
+from aws_advanced_python_wrapper.utils.properties import Properties, WrapperProperties
 from aws_advanced_python_wrapper.utils.utils import Utils
 
 logger = Logger(__name__)
@@ -52,7 +51,8 @@ class DriverDialectManager(DriverDialectProvider):
     }
 
     pool_connection_driver_dialect: Dict[str, str] = {
-        "SqlAlchemyPooledConnectionProvider": "aws_advanced_python_wrapper.sqlalchemy_driver_dialect.SqlAlchemyDriverDialect"
+        "SqlAlchemyPooledConnectionProvider": "aws_advanced_python_wrapper.sqlalchemy_driver_dialect.SqlAlchemyDriverDialect",
+        "SqlAlchemyTortoisePooledConnectionProvider": "aws_advanced_python_wrapper.sqlalchemy_driver_dialect.SqlAlchemyDriverDialect",
     }
 
     @staticmethod

aws_advanced_python_wrapper/mysql_driver_dialect.py

@@ -19,6 +19,7 @@
 if TYPE_CHECKING:
     from aws_advanced_python_wrapper.hostinfo import HostInfo
     from aws_advanced_python_wrapper.pep249 import Connection
+    from types import ModuleType
 
 from concurrent.futures import Executor, ThreadPoolExecutor, TimeoutError
 from inspect import signature
@@ -28,12 +29,11 @@
 from aws_advanced_python_wrapper.errors import UnsupportedOperationError
 from aws_advanced_python_wrapper.utils.decorators import timeout
 from aws_advanced_python_wrapper.utils.messages import Messages
-from aws_advanced_python_wrapper.utils.properties import (Properties,
-                                                          PropertiesUtils,
-                                                          WrapperProperties)
+from aws_advanced_python_wrapper.utils.properties import Properties, PropertiesUtils, WrapperProperties
 
 CMYSQL_ENABLED = False
 
+import mysql.connector
 from mysql.connector import MySQLConnection  # noqa: E402
 from mysql.connector.cursor import MySQLCursor  # noqa: E402
 
@@ -201,3 +201,6 @@ def prepare_connect_info(self, host_info: HostInfo, original_props: Properties)
 
     def supports_connect_timeout(self) -> bool:
         return True
+
+    def get_driver_module(self) -> ModuleType:
+        return mysql.connector

aws_advanced_python_wrapper/pg_driver_dialect.py

@@ -14,23 +14,21 @@
 
 from __future__ import annotations
 
+from inspect import signature
 from typing import TYPE_CHECKING, Any, Callable, Set
 
 import psycopg
 
 if TYPE_CHECKING:
     from aws_advanced_python_wrapper.hostinfo import HostInfo
     from aws_advanced_python_wrapper.pep249 import Connection
-
-from inspect import signature
+    from types import ModuleType
 
 from aws_advanced_python_wrapper.driver_dialect import DriverDialect
 from aws_advanced_python_wrapper.driver_dialect_codes import DriverDialectCodes
 from aws_advanced_python_wrapper.errors import UnsupportedOperationError
 from aws_advanced_python_wrapper.utils.messages import Messages
-from aws_advanced_python_wrapper.utils.properties import (Properties,
-                                                          PropertiesUtils,
-                                                          WrapperProperties)
+from aws_advanced_python_wrapper.utils.properties import Properties, PropertiesUtils, WrapperProperties
 
 
 class PgDriverDialect(DriverDialect):
@@ -175,3 +173,6 @@ def supports_tcp_keepalive(self) -> bool:
 
     def supports_abort_connection(self) -> bool:
         return True
+
+    def get_driver_module(self) -> ModuleType:
+        return psycopg

aws_advanced_python_wrapper/resources/aws_advanced_python_wrapper_messages.properties

@@ -385,6 +385,8 @@ WeightedRandomHostSelector.WeightedRandomInvalidDefaultWeight=[WeightedRandomHos
 SqlAlchemyPooledConnectionProvider.PoolNone=[SqlAlchemyPooledConnectionProvider] Attempted to find or create a pool for '{}' but the result of the attempt evaluated to None.
 SqlAlchemyPooledConnectionProvider.UnableToCreateDefaultKey=[SqlAlchemyPooledConnectionProvider] Unable to create a default key for internal connection pools. By default, the user parameter is used, but the given user evaluated to None or the empty string (""). Please ensure you have passed a valid user in the connection properties.
 
+SqlAlchemyTortoiseConnectionProvider.UnableToDetermineDialect=[SqlAlchemyTortoiseConnectionProvider] Unable to resolve sql alchemy dialect for the following driver dialect '{}'.
+
 SqlAlchemyDriverDialect.SetValueOnNoneConnection=[SqlAlchemyDriverDialect] Attempted to set the '{}' value on a pooled connection, but no underlying driver connection was found. This can happen if the pooled connection has previously been closed.
 
 StaleDnsHelper.ClusterEndpointDns=[StaleDnsPlugin] Cluster endpoint {} resolves to {}.

aws_advanced_python_wrapper/sqlalchemy_driver_dialect.py

@@ -16,17 +16,18 @@
 
 from typing import TYPE_CHECKING, Any
 
-from aws_advanced_python_wrapper.driver_dialect import DriverDialect
-from aws_advanced_python_wrapper.errors import AwsWrapperError
-from aws_advanced_python_wrapper.utils.messages import Messages
-
 if TYPE_CHECKING:
     from aws_advanced_python_wrapper.hostinfo import HostInfo
     from aws_advanced_python_wrapper.pep249 import Connection
     from aws_advanced_python_wrapper.utils.properties import Properties
+    from types import ModuleType
 
 from sqlalchemy import PoolProxiedConnection
 
+from aws_advanced_python_wrapper.driver_dialect import DriverDialect
+from aws_advanced_python_wrapper.errors import AwsWrapperError
+from aws_advanced_python_wrapper.utils.messages import Messages
+
 
 class SqlAlchemyDriverDialect(DriverDialect):
     _driver_name: str = "SQLAlchemy"
@@ -125,3 +126,6 @@ def transfer_session_state(self, from_conn: Connection, to_conn: Connection):
             return None
 
         return self._underlying_driver.transfer_session_state(from_driver_conn, to_driver_conn)
+    
+    def get_driver_module(self) -> ModuleType:
+        return self._underlying_driver.get_driver_module()

aws_advanced_python_wrapper/tortoise/backend/mysql/client.py

@@ -22,7 +22,6 @@
 from mysql.connector import errors
 from mysql.connector.charsets import MYSQL_CHARACTER_SETS
 from pypika_tortoise import MySQLQuery
-from tortoise import timezone
 from tortoise.backends.base.client import (
     BaseDBAsyncClient,
     Capabilities,
@@ -39,16 +38,16 @@
 )
 
 from aws_advanced_python_wrapper.connection_provider import ConnectionProviderManager
+from aws_advanced_python_wrapper.errors import AwsWrapperError, FailoverError
 from aws_advanced_python_wrapper.hostinfo import HostInfo
-from aws_advanced_python_wrapper.sql_alchemy_connection_provider import SqlAlchemyPooledConnectionProvider
 from aws_advanced_python_wrapper.tortoise.backend.base.client import (
     TortoiseAwsClientConnectionWrapper,
     TortoiseAwsClientTransactionContext,
 )
 from aws_advanced_python_wrapper.tortoise.backend.mysql.executor import AwsMySQLExecutor
 from aws_advanced_python_wrapper.tortoise.backend.mysql.schema_generator import AwsMySQLSchemaGenerator
+from aws_advanced_python_wrapper.tortoise.sql_alchemy_tortoise_connection_provider import SqlAlchemyTortoisePooledConnectionProvider
 from aws_advanced_python_wrapper.utils.log import Logger
-from aws_advanced_python_wrapper.errors import AwsWrapperError
 
 logger = Logger(__name__)
 T = TypeVar("T")
@@ -62,10 +61,12 @@ async def translate_exceptions_(self, *args) -> T:
         try:
             try:
                 return await func(self, *args)
-            except AwsWrapperError as aws_err:
+            except AwsWrapperError as aws_err: # Unwrap any AwsWrappedErrors
                 if aws_err.__cause__:
                     raise aws_err.__cause__
                 raise
+        except FailoverError as exc: # Raise any failover errors
+            raise
         except errors.IntegrityError as exc:
             raise IntegrityError(exc)
         except (
@@ -127,16 +128,16 @@ def __init__(
         self.extra.pop("connection_name", None)
         self.extra.pop("fetch_inserted", None)
         self.extra.pop("db", None)
-        self.extra.pop("autocommit", None)
+        self.extra.pop("autocommit", None) # We need this to be true
         self.extra.setdefault("sql_mode", "STRICT_TRANS_TABLES")
 
         # Initialize connection templates
         self._init_connection_templates()
-        
+
         # Initialize state
         self._template = {}
         self._connection = None
-        self._pool = None
+        self._pool: SqlAlchemyTortoisePooledConnectionProvider = None
         self._pool_init_lock = asyncio.Lock()
 
     def _init_connection_templates(self) -> None:
@@ -158,20 +159,19 @@ def _configure_pool(self, host_info: HostInfo, props: Dict[str, Any]) -> Dict[st
         """Configure connection pool settings."""
         return {"pool_size": self.pool_maxsize, "max_overflow": -1}
 
-    @staticmethod
-    def _get_pool_key(host_info: HostInfo, props: Dict[str, Any]) -> str:
+    def _get_pool_key(self, host_info: HostInfo, props: Dict[str, Any]) -> str:
         """Generate unique pool key for connection pooling."""
         url = host_info.url
         user = props["user"]
         db = props["database"]
-        return f"{url}{user}{db}"
+        return f"{url}{user}{db}{self.connection_name}"
 
     async def _init_pool_if_needed(self) -> None:
         """Initialize connection pool only once across all instances."""
         if not AwsMySQLClient._pool_initialized:
             async with AwsMySQLClient._pool_init_class_lock:
                 if not AwsMySQLClient._pool_initialized:
-                    self._pool = SqlAlchemyPooledConnectionProvider(
+                    self._pool = SqlAlchemyTortoisePooledConnectionProvider(
                         pool_configurator=self._configure_pool,
                         pool_mapping=self._get_pool_key,
                     )

aws_advanced_python_wrapper/tortoise/sql_alchemy_tortoise_connection_provider.py

@@ -0,0 +1,85 @@
+#  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 types import ModuleType
+from typing import Callable, Dict
+
+from sqlalchemy import Dialect
+from sqlalchemy.dialects.mysql import mysqlconnector
+from sqlalchemy.dialects.postgresql import psycopg
+
+from aws_advanced_python_wrapper.database_dialect import DatabaseDialect
+from aws_advanced_python_wrapper.driver_dialect import DriverDialect
+from aws_advanced_python_wrapper.errors import AwsWrapperError
+from aws_advanced_python_wrapper.hostinfo import HostInfo
+from aws_advanced_python_wrapper.sql_alchemy_connection_provider import SqlAlchemyPooledConnectionProvider
+from aws_advanced_python_wrapper.utils.log import Logger
+from aws_advanced_python_wrapper.utils.messages import Messages
+from aws_advanced_python_wrapper.utils.properties import Properties
+
+logger = Logger(__name__)
+
+
+
+class SqlAlchemyTortoisePooledConnectionProvider(SqlAlchemyPooledConnectionProvider):
+    """
+    Tortoise-specific pooled connection provider that handles failover by disposing pools.
+    """
+
+    _sqlalchemy_dialect_map : Dict[str, ModuleType] = {
+        "MySQLDriverDialect": mysqlconnector,
+        "PostgresDriverDialect": psycopg
+    }
+
+    def accepts_host_info(self, host_info: HostInfo, props: Properties) -> bool:
+        if self._accept_url_func:
+            return self._accept_url_func(host_info, props)
+        url_type = SqlAlchemyPooledConnectionProvider._rds_utils.identify_rds_type(host_info.host)
+        return url_type.is_rds
+
+    def _create_pool(
+            self,
+            target_func: Callable,
+            driver_dialect: DriverDialect,
+            database_dialect: DatabaseDialect,
+            host_info: HostInfo,
+            props: Properties):
+        kwargs = dict() if self._pool_configurator is None else self._pool_configurator(host_info, props)
+        prepared_properties = driver_dialect.prepare_connect_info(host_info, props)
+        database_dialect.prepare_conn_props(prepared_properties)
+        kwargs["creator"] = self._get_connection_func(target_func, prepared_properties)
+        dialect = self._get_pool_dialect(driver_dialect)
+        if not dialect:
+            raise AwsWrapperError(Messages.get_formatted("SqlAlchemyTortoisePooledConnectionProvider.NoDialect", driver_dialect.__class__.__name__))
+
+        '''
+        We need to pass in pre_ping and dialect to QueuePool the queue pool to enable health checks.
+        Without this health check, we could be using dead connections after a failover.
+        '''
+        kwargs["pre_ping"] = True
+        kwargs["dialect"] = dialect
+        return self._create_sql_alchemy_pool(**kwargs)
+    
+    def _get_pool_dialect(self, driver_dialect: DriverDialect) -> Dialect:
+        dialect = None
+        driver_dialect_class_name = driver_dialect.__class__.__name__
+        if driver_dialect_class_name == "SqlAlchemyDriverDialect":
+            driver_dialect_class_name = driver_dialect_class_name._underlying_driver_dialect.__class__.__name__
+        module = self._sqlalchemy_dialect_map.get(driver_dialect_class_name)
+
+        if not module:
+            return dialect
+        dialect = module.dialect()
+        dialect.dbapi = driver_dialect.get_driver_module()
+
+        return dialect

tests/integration/container/utils/test_telemetry_info.py

@@ -11,18 +11,7 @@
 #  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.
-#
-#  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.
+
 import typing
 from typing import Any, Dict