feat: add Tortoise ORM support

.github/workflows/integration_tests.yml

@@ -5,6 +5,7 @@ on:
   push:
     branches:
       - main
+      - feat/tortoise
 
 permissions:
   id-token: write   # This is required for requesting the JWT
@@ -17,8 +18,9 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: [ "3.11", "3.12", "3.13"  ]
-        environment: [ "mysql", "pg" ]
+        python-version: [ "3.11", "3.12", "3.13"]
+        # environment: [ "mysql", "pg" ]
+        environment: [ "mysql",]
 
     steps:
       - name: 'Clone repository'
@@ -48,8 +50,10 @@ jobs:
           aws-region: ${{ secrets.AWS_DEFAULT_REGION }}
 
       - name: 'Run LTS Integration Tests'
+        timeout-minutes: 120
         run: |
-          ./gradlew --no-parallel --no-daemon test-python-${{ matrix.python-version }}-${{ matrix.environment }} --info
+          ./gradlew --no-parallel --no-daemon test-python-${{ matrix.python-version }}-${{ matrix.environment }} --info \
+          2>&1 | tee test-output.log
         env:
           RDS_CLUSTER_DOMAIN: ${{ secrets.DB_CONN_SUFFIX }}
           RDS_DB_REGION: ${{ secrets.AWS_DEFAULT_REGION }}
@@ -64,6 +68,13 @@ jobs:
           path: ./tests/integration/container/reports
           retention-days: 5
 
+      - name: Upload logs
+        uses: actions/upload-artifact@v4
+        if: always()
+        with:
+          name: pytest-integration-${{ matrix.python-version }}-${{ matrix.environment }}-log
+          path: test-output.log
+
   latest-integration-tests:
     name: Run Latest Integration Tests
     runs-on: ubuntu-latest

aws_advanced_python_wrapper/custom_endpoint_plugin.py

@@ -169,7 +169,8 @@ def _run(self):
                             len(endpoints),
                             endpoint_hostnames)
 
-                        sleep(self._refresh_rate_ns / 1_000_000_000)
+                        if self._stop_event.wait(self._refresh_rate_ns / 1_000_000_000):
+                            break
                         continue
 
                     endpoint_info = CustomEndpointInfo.from_db_cluster_endpoint(endpoints[0])
@@ -178,7 +179,8 @@ def _run(self):
                     if cached_info is not None and cached_info == endpoint_info:
                         elapsed_time = perf_counter_ns() - start_ns
                         sleep_duration = max(0, self._refresh_rate_ns - elapsed_time)
-                        sleep(sleep_duration / 1_000_000_000)
+                        if self._stop_event.wait(sleep_duration / 1_000_000_000):
+                            break
                         continue
 
                     logger.debug(
@@ -196,7 +198,8 @@ def _run(self):
 
                     elapsed_time = perf_counter_ns() - start_ns
                     sleep_duration = max(0, self._refresh_rate_ns - elapsed_time)
-                    sleep(sleep_duration / 1_000_000_000)
+                    if self._stop_event.wait(sleep_duration / 1_000_000_000):
+                        break
                     continue
                 except InterruptedError as e:
                     raise e

aws_advanced_python_wrapper/errors.py

@@ -45,3 +45,31 @@ class FailoverSuccessError(FailoverError):
 
 class ReadWriteSplittingError(AwsWrapperError):
     __module__ = "aws_advanced_python_wrapper"
+
+
+class AsyncConnectionPoolError(AwsWrapperError):
+    __module__ = "aws_advanced_python_wrapper"
+
+
+class PoolNotInitializedError(AsyncConnectionPoolError):
+    __module__ = "aws_advanced_python_wrapper"
+
+
+class PoolClosingError(AsyncConnectionPoolError):
+    __module__ = "aws_advanced_python_wrapper"
+
+
+class PoolExhaustedError(AsyncConnectionPoolError):
+    __module__ = "aws_advanced_python_wrapper"
+
+
+class ConnectionReleasedError(AsyncConnectionPoolError):
+    __module__ = "aws_advanced_python_wrapper"
+
+
+class PoolSizeLimitError(AsyncConnectionPoolError):
+    __module__ = "aws_advanced_python_wrapper"
+
+
+class PoolHealthCheckError(AsyncConnectionPoolError):
+    __module__ = "aws_advanced_python_wrapper"

aws_advanced_python_wrapper/iam_plugin.py

@@ -114,7 +114,7 @@ def _connect(self, host_info: HostInfo, props: Properties, connect_func: Callabl
 
             is_cached_token = (token_info is not None and not token_info.is_expired())
             if not self._plugin_service.is_login_exception(error=e) or not is_cached_token:
-                raise AwsWrapperError(Messages.get_formatted("IamAuthPlugin.ConnectException", e)) from e
+                raise
 
             # Login unsuccessful with cached token
             # Try to generate a new token and try to connect again

aws_advanced_python_wrapper/plugin_service.py

@@ -703,8 +703,7 @@ def is_plugin_in_use(self, plugin_class: Type[Plugin]) -> bool:
 
     def release_resources(self):
         try:
-            if self.current_connection is not None and not self.driver_dialect.is_closed(
-                    self.current_connection):
+            if self.current_connection is not None:
                 self.current_connection.close()
         except Exception:
             # ignore
@@ -845,6 +844,9 @@ def get_plugins(self) -> List[Plugin]:
             if plugin_codes is None:
                 plugin_codes = WrapperProperties.DEFAULT_PLUGINS
 
+            if plugin_codes.lower() == "none":
+                plugin_codes = ""
+
             if plugin_codes != "":
                 plugin_factories = self.create_plugin_factories_from_list(plugin_codes.split(","))
 

aws_advanced_python_wrapper/tortoise_orm/__init__.py

@@ -0,0 +1,47 @@
+#  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 tortoise.backends.base.config_generator import DB_LOOKUP
+
+
+def cast_to_bool(value):
+    """Generic function to cast various types to boolean."""
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, str):
+        return value.lower() in ('true', '1', 'yes', 'on')
+    return bool(value)
+
+
+# Register AWS MySQL backend
+DB_LOOKUP["aws-mysql"] = {
+    "engine": "aws_advanced_python_wrapper.tortoise_orm.backends.mysql",
+    "vmap": {
+        "path": "database",
+        "hostname": "host",
+        "port": "port",
+        "username": "user",
+        "password": "password",
+    },
+    "defaults": {"port": 3306, "charset": "utf8mb4", "sql_mode": "STRICT_TRANS_TABLES"},
+    "cast": {
+        "minsize": int,
+        "maxsize": int,
+        "connect_timeout": int,
+        "echo": cast_to_bool,
+        "use_unicode": cast_to_bool,
+        "ssl": cast_to_bool,
+        "use_pure": cast_to_bool
+    },
+}

aws_advanced_python_wrapper/tortoise_orm/async_support/__init__.py

@@ -0,0 +1,13 @@
+#  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.