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

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/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.