fix static check

Author: faviansamatha

aws_advanced_python_wrapper/driver_dialect.py

@@ -165,7 +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,7 +24,8 @@
 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__)

aws_advanced_python_wrapper/mysql_driver_dialect.py

@@ -11,11 +11,12 @@
 #  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 typing import TYPE_CHECKING, Any, Callable, ClassVar, Set
 
+import mysql.connector
+
 if TYPE_CHECKING:
     from aws_advanced_python_wrapper.hostinfo import HostInfo
     from aws_advanced_python_wrapper.pep249 import Connection
@@ -29,11 +30,12 @@
 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
 

aws_advanced_python_wrapper/pg_driver_dialect.py

@@ -28,7 +28,9 @@
 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):

aws_advanced_python_wrapper/sql_alchemy_connection_provider.py

@@ -88,7 +88,7 @@ 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 RdsUrlType.RDS_INSTANCE == url_type or RdsUrlType.RDS_WRITER_CLUSTER
+        return url_type in (RdsUrlType.RDS_INSTANCE, RdsUrlType.RDS_WRITER_CLUSTER)
 
     def accepts_strategy(self, role: HostRole, strategy: str) -> bool:
         return strategy == SqlAlchemyPooledConnectionProvider._LEAST_CONNECTIONS or strategy in self._accepted_strategies

aws_advanced_python_wrapper/sqlalchemy_driver_dialect.py

@@ -32,6 +32,7 @@
 class SqlAlchemyDriverDialect(DriverDialect):
     _driver_name: str = "SQLAlchemy"
     TARGET_DRIVER_CODE: str = "sqlalchemy"
+    _underlying_driver_dialect = None
 
     def __init__(self, underlying_driver: DriverDialect, props: Properties):
         super().__init__(props)
@@ -126,6 +127,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/__init__.py

@@ -18,12 +18,12 @@
 DB_LOOKUP["aws-mysql"] = {
     "engine": "aws_advanced_python_wrapper.tortoise.backend.mysql",
     "vmap": {
-            "path": "database",
-            "hostname": "host",
-            "port": "port",
-            "username": "user",
-            "password": "password",
-        },
+        "path": "database",
+        "hostname": "host",
+        "port": "port",
+        "username": "user",
+        "password": "password",
+    },
     "defaults": {"port": 3306, "charset": "utf8mb4", "sql_mode": "STRICT_TRANS_TABLES"},
     "cast": {
         "minsize": int,

aws_advanced_python_wrapper/tortoise/backend/base/client.py

@@ -12,12 +12,16 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+from __future__ import annotations
+
 import asyncio
-import mysql.connector
 from contextlib import asynccontextmanager
-from typing import Any, Callable, Generic
+from typing import Any, Callable, Dict, Generic, cast
 
-from tortoise.backends.base.client import BaseDBAsyncClient, T_conn, TransactionalDBClient, TransactionContext
+import mysql.connector
+from tortoise.backends.base.client import (BaseDBAsyncClient, T_conn,
+                                           TransactionalDBClient,
+                                           TransactionContext)
 from tortoise.connection import connections
 from tortoise.exceptions import TransactionManagementError
 
@@ -26,55 +30,55 @@
 
 class AwsWrapperAsyncConnector:
     """Class for creating and closing AWS wrapper connections."""
-    
+
     @staticmethod
-    async def ConnectWithAwsWrapper(connect_func: Callable, **kwargs) -> AwsWrapperConnection:
+    async def connect_with_aws_wrapper(connect_func: Callable, **kwargs) -> AwsConnectionAsyncWrapper:
         """Create an AWS wrapper connection with async cursor support."""
         connection = await asyncio.to_thread(
             AwsWrapperConnection.connect, connect_func, **kwargs
         )
         return AwsConnectionAsyncWrapper(connection)
-    
+
     @staticmethod
-    async def CloseAwsWrapper(connection: AwsWrapperConnection) -> None:
+    async def close_aws_wrapper(connection: AwsWrapperConnection) -> None:
         """Close an AWS wrapper connection asynchronously."""
         await asyncio.to_thread(connection.close)
 
 
 class AwsCursorAsyncWrapper:
     """Wraps sync AwsCursor cursor with async support."""
-    
+
     def __init__(self, sync_cursor):
         self._cursor = sync_cursor
-    
+
     async def execute(self, query, params=None):
         """Execute a query asynchronously."""
         return await asyncio.to_thread(self._cursor.execute, query, params)
-    
+
     async def executemany(self, query, params_list):
         """Execute multiple queries asynchronously."""
         return await asyncio.to_thread(self._cursor.executemany, query, params_list)
-    
+
     async def fetchall(self):
         """Fetch all results asynchronously."""
         return await asyncio.to_thread(self._cursor.fetchall)
-    
+
     async def fetchone(self):
         """Fetch one result asynchronously."""
         return await asyncio.to_thread(self._cursor.fetchone)
-    
+
     async def close(self):
         """Close cursor asynchronously."""
         return await asyncio.to_thread(self._cursor.close)
-    
+
     def __getattr__(self, name):
         """Delegate non-async attributes to the wrapped cursor."""
         return getattr(self._cursor, name)
 
 
 class AwsConnectionAsyncWrapper(AwsWrapperConnection):
     """Wraps sync AwsConnection with async cursor support."""
-    
+
     def __init__(self, connection: AwsWrapperConnection):
         self._wrapped_connection = connection
 
@@ -90,40 +94,50 @@ async def cursor(self):
     async def rollback(self):
         """Rollback the current transaction."""
         return await asyncio.to_thread(self._wrapped_connection.rollback)
-    
+
     async def commit(self):
         """Commit the current transaction."""
         return await asyncio.to_thread(self._wrapped_connection.commit)
-    
+
     async def set_autocommit(self, value: bool):
         """Set autocommit mode."""
         return await asyncio.to_thread(setattr, self._wrapped_connection, 'autocommit', value)
 
     def __getattr__(self, name):
         """Delegate all other attributes/methods to the wrapped connection."""
         return getattr(self._wrapped_connection, name)
-    
+
     def __del__(self):
         """Delegate cleanup to wrapped connection."""
         if hasattr(self, '_wrapped_connection'):
             # Let the wrapped connection handle its own cleanup
             pass
 
 
+class AwsBaseDBAsyncClient(BaseDBAsyncClient):
+    _template: Dict[str, Any]
+
+
+class AwsTransactionalDBClient(TransactionalDBClient):
+    _template: Dict[str, Any]
+    _parent: AwsBaseDBAsyncClient
+    pass
+
+
 class TortoiseAwsClientConnectionWrapper(Generic[T_conn]):
     """Manages acquiring from and releasing connections to a pool."""
 
     __slots__ = ("client", "connection", "connect_func", "with_db")
 
     def __init__(
-        self, 
-        client: BaseDBAsyncClient, 
-        connect_func: Callable, 
+        self,
+        client: AwsBaseDBAsyncClient,
+        connect_func: Callable,
         with_db: bool = True
     ) -> None:
         self.connect_func = connect_func
         self.client = client
-        self.connection: T_conn | None = None
+        self.connection: AwsConnectionAsyncWrapper | None = None
         self.with_db = with_db
 
     async def ensure_connection(self) -> None:
@@ -133,22 +147,22 @@ async def ensure_connection(self) -> None:
     async def __aenter__(self) -> T_conn:
         """Acquire connection from pool."""
         await self.ensure_connection()
-        self.connection = await AwsWrapperAsyncConnector.ConnectWithAwsWrapper(self.connect_func, **self.client._template)
-        return self.connection
+        self.connection = await AwsWrapperAsyncConnector.connect_with_aws_wrapper(self.connect_func, **self.client._template)
+        return cast("T_conn", self.connection)
 
     async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         """Close connection and release back to pool."""
         if self.connection:
-            await AwsWrapperAsyncConnector.CloseAwsWrapper(self.connection)
+            await AwsWrapperAsyncConnector.close_aws_wrapper(self.connection)
 
 
 class TortoiseAwsClientTransactionContext(TransactionContext):
     """Transaction context that uses a pool to acquire connections."""
 
     __slots__ = ("client", "connection_name", "token")
 
-    def __init__(self, client: TransactionalDBClient) -> None:
-        self.client = client
+    def __init__(self, client: AwsTransactionalDBClient) -> None:
+        self.client: AwsTransactionalDBClient = client
         self.connection_name = client.connection_name
 
     async def ensure_connection(self) -> None:
@@ -158,13 +172,13 @@ async def ensure_connection(self) -> None:
     async def __aenter__(self) -> TransactionalDBClient:
         """Enter transaction context."""
         await self.ensure_connection()
-        
+
         # Set the context variable so the current task sees a TransactionWrapper connection
         self.token = connections.set(self.connection_name, self.client)
-        
+
         # Create connection and begin transaction
-        self.client._connection = await AwsWrapperAsyncConnector.ConnectWithAwsWrapper(
-            mysql.connector.Connect, 
+        self.client._connection = await AwsWrapperAsyncConnector.connect_with_aws_wrapper(
+            mysql.connector.Connect,
             **self.client._parent._template
         )
         await self.client.begin()
@@ -181,5 +195,5 @@ async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
                 else:
                     await self.client.commit()
         finally:
-            await AwsWrapperAsyncConnector.CloseAwsWrapper(self.client._connection)
+            await AwsWrapperAsyncConnector.close_aws_wrapper(self.client._connection)
             connections.reset(self.token)