Add api changed from sqlalchemy/sqlalchemy#6832 (#177)

Author: CaselIT

sqlalchemy-stubs/dialects/mysql/aiomysql.pyi

@@ -2,6 +2,7 @@ from typing import Any
 from typing import Optional
 
 from .pymysql import MySQLDialect_pymysql as MySQLDialect_pymysql
+from ...engine import AdaptedConnection
 from ... import pool as pool
 from ... import util as util
 from ...util.concurrency import await_fallback as await_fallback
@@ -39,7 +40,7 @@ class AsyncAdapt_aiomysql_ss_cursor(AsyncAdapt_aiomysql_cursor):
     def fetchmany(self, size: Optional[Any] = ...): ...
     def fetchall(self): ...
 
-class AsyncAdapt_aiomysql_connection:
+class AsyncAdapt_aiomysql_connection(AdaptedConnection):
     await_: Any = ...
     dbapi: Any = ...
     def __init__(self, dbapi: Any, connection: Any) -> None: ...
@@ -71,5 +72,6 @@ class MySQLDialect_aiomysql(MySQLDialect_pymysql):
     def get_pool_class(cls, url: Any): ...
     def create_connect_args(self, url: Any): ...
     def is_disconnect(self, e: Any, connection: Any, cursor: Any): ...
+    def get_driver_connection(self, connection: Any) -> Any:
 
 dialect = MySQLDialect_aiomysql

sqlalchemy-stubs/dialects/postgresql/asyncpg.pyi

@@ -11,6 +11,7 @@ from .base import PGExecutionContext as PGExecutionContext
 from .base import PGIdentifierPreparer as PGIdentifierPreparer
 from .base import REGCLASS as REGCLASS
 from .base import UUID as UUID
+from ...engine import AdaptedConnection
 from ... import exc as exc
 from ... import pool as pool
 from ... import processors as processors
@@ -117,7 +118,7 @@ class AsyncAdapt_asyncpg_ss_cursor(AsyncAdapt_asyncpg_cursor):
     def fetchall(self): ...
     def executemany(self, operation: Any, seq_of_parameters: Any) -> None: ...
 
-class AsyncAdapt_asyncpg_connection:
+class AsyncAdapt_asyncpg_connection(AdaptedConnection):
     await_: Any = ...
     dbapi: Any = ...
     isolation_level: str = ...
@@ -211,5 +212,6 @@ class PGDialect_asyncpg(PGDialect):
         self, cursor: Any, list_of_tuples: Any, context: Any
     ) -> None: ...
     def on_connect(self): ...
+    def get_driver_connection(self, connection: Any) -> Any:
 
 dialect = PGDialect_asyncpg

sqlalchemy-stubs/engine/__init__.pyi

@@ -15,6 +15,7 @@ from .cursor import BufferedRowResultProxy as BufferedRowResultProxy
 from .cursor import CursorResult as CursorResult
 from .cursor import FullyBufferedResultProxy as FullyBufferedResultProxy
 from .cursor import LegacyCursorResult as LegacyCursorResult
+from .interfaces import AdaptedConnection as AdaptedConnection
 from .interfaces import Compiled as Compiled
 from .interfaces import Connectable as Connectable
 from .interfaces import CreateEnginePlugin as CreateEnginePlugin

sqlalchemy-stubs/engine/default.pyi

@@ -198,6 +198,9 @@ class DefaultDialect(interfaces.Dialect):
     ) -> None: ...
     def normalize_name(self, name: Optional[str]) -> Optional[str]: ...
     def denormalize_name(self, name: Optional[str]) -> Optional[str]: ...
+    def get_driver_connection(
+        self, connection: interfaces._DBAPIConnection
+    ) -> Any: ...
 
 class StrCompileDialect(DefaultDialect):
     statement_compiler: Type[compiler.StrSQLCompiler] = ...

sqlalchemy-stubs/engine/interfaces.pyi

@@ -300,6 +300,7 @@ class Dialect:
     def load_provisioning(cls) -> None: ...
     @classmethod
     def engine_created(cls, engine: Engine) -> None: ...
+    def get_driver_connection(self, connection: _DBAPIConnection) -> Any: ...
 
 class CreateEnginePlugin:
     url: URL = ...
@@ -362,3 +363,7 @@ class ExceptionContext:
     execution_context: Optional[ExecutionContext] = ...
     is_disconnect: bool = ...
     invalidate_pool_on_disconnect: bool = ...
+
+class AdaptedConnection:
+    @property
+    def driver_connection(self) -> Any: ...

sqlalchemy-stubs/pool/base.pyi

@@ -1,6 +1,7 @@
 from typing import Any
 from typing import Optional
 
+from ..engine.interfaces import _DBAPIConnection
 from .. import event as event
 from .. import exc as exc
 from .. import log as log
@@ -43,7 +44,13 @@ class _ConnectionRecord:
     fresh: bool = ...
     fairy_ref: Any = ...
     starttime: Any = ...
-    connection: Any = ...
+    dbapi_connection: _DBAPIConnection = ...
+    @property
+    def driver_connection(self) -> Any: ...
+    @property
+    def connection(self) -> _DBAPIConnection: ...
+    @connection.setter
+    def connection(self, value: _DBAPIConnection) -> None: ...
     @util.memoized_property
     def info(self): ...
     @util.memoized_property
@@ -60,7 +67,13 @@ class _ConnectionRecord:
     def get_connection(self): ...
 
 class _ConnectionFairy:
-    connection: Any = ...
+    dbapi_connection: _DBAPIConnection = ...
+    @property
+    def driver_connection(self) -> Any: ...
+    @property
+    def connection(self) -> _DBAPIConnection: ...
+    @connection.setter
+    def connection(self, value: _DBAPIConnection) -> None: ...
     def __init__(
         self, dbapi_connection: Any, connection_record: Any, echo: Any
     ) -> None: ...