fix: allow async session factory

.github/workflows/tests.yml

@@ -11,7 +11,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.7", "3.8", "3.9", "3.10"]
+        python-version: ["3.7", "3.8", "3.9", "3.10", "3.11", "3.12"]
     steps:
     - uses: actions/checkout@v2
     - name: Set up Kerberos
@@ -30,12 +30,15 @@ jobs:
       run: |
         make test
     - name: Lint
+      if: matrix.python-version != '3.7'
       run: |
         make lint
     - name: Format
+      if: matrix.python-version != '3.7'
       run: |
         make format-check
     - name: Type annotations
+      if: matrix.python-version != '3.7'
       run: |
         make types
   mysql-connector-j:

.gitignore

@@ -106,6 +106,7 @@ celerybeat.pid
 .venv
 env/
 venv/
+venv*/
 ENV/
 env.bak/
 venv.bak/

.pylintrc

@@ -15,6 +15,7 @@ disable=
     unused-argument,
     redefined-outer-name,
     too-many-statements,
+    multiple-statements,
 
 [FORMAT]
 

integration/run.py

@@ -72,9 +72,11 @@ async def get_user(self, username):
                 password = user.get("password")
                 return User(
                     name=username,
-                    auth_string=NativePasswordAuthPlugin.create_auth_string(password)
-                    if password
-                    else None,
+                    auth_string=(
+                        NativePasswordAuthPlugin.create_auth_string(password)
+                        if password
+                        else None
+                    ),
                     auth_plugin=NativePasswordAuthPlugin.name,
                 )
             elif auth_plugin == "authentication_kerberos":

mysql_mimic/__init__.py

@@ -1,4 +1,5 @@
 """Implementation of the mysql server wire protocol"""
+
 from mysql_mimic.auth import (
     User,
     IdentityProvider,

mysql_mimic/connection.py

@@ -544,9 +544,11 @@ async def handle_stmt_fetch(self, data: bytes) -> None:
 
         await self.stream.write(
             self.ok_or_eof(
-                flags=types.ServerStatus.SERVER_STATUS_LAST_ROW_SENT
-                if done
-                else types.ServerStatus.SERVER_STATUS_CURSOR_EXISTS
+                flags=(
+                    types.ServerStatus.SERVER_STATUS_LAST_ROW_SENT
+                    if done
+                    else types.ServerStatus.SERVER_STATUS_CURSOR_EXISTS
+                )
             )
         )
 

mysql_mimic/schema.py

@@ -275,8 +275,7 @@ class BaseInfoSchema:
     Base InfoSchema interface used by the `Session` class.
     """
 
-    async def query(self, expression: exp.Expression) -> AllowedResult:
-        ...
+    async def query(self, expression: exp.Expression) -> AllowedResult: ...
 
 
 class InfoSchema(BaseInfoSchema):

mysql_mimic/server.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
 import asyncio
+import inspect
 from ssl import SSLContext
 from socket import socket
-from typing import Callable, Any, Optional, Sequence
+from typing import Callable, Any, Optional, Sequence, Awaitable
 
 from mysql_mimic.auth import IdentityProvider, SimpleIdentityProvider
 from mysql_mimic.connection import Connection
@@ -36,7 +37,7 @@ class MysqlServer:
 
     def __init__(
         self,
-        session_factory: Callable[[], BaseSession] = Session,
+        session_factory: Callable[[], BaseSession | Awaitable[BaseSession]] = Session,
         capabilities: Capabilities = DEFAULT_SERVER_CAPABILITIES,
         control: Control | None = None,
         identity_provider: IdentityProvider | None = None,
@@ -57,9 +58,14 @@ async def _client_connected_cb(
     ) -> None:
         stream = MysqlStream(reader, writer)
 
+        if inspect.iscoroutinefunction(self.session_factory):
+            session = await self.session_factory()
+        else:
+            session = self.session_factory()
+
         connection = Connection(
             stream=stream,
-            session=self.session_factory(),
+            session=session,
             control=self.control,
             server_capabilities=self.capabilities,
             identity_provider=self.identity_provider,

mysql_mimic/utils.py

@@ -55,7 +55,9 @@ def nonce(nbytes: int) -> bytes:
 
 def find_tables(expression: exp.Expression) -> List[exp.Table]:
     """Find all tables in an expression"""
-    if isinstance(expression, (exp.Subqueryable, exp.Subquery)):
+    if isinstance(
+        expression, (exp.Select, exp.Subquery, exp.Union, exp.Except, exp.Intersect)
+    ):
         return [
             source
             for scope in traverse_scope(expression)

mysql_mimic/variables.py

@@ -10,8 +10,7 @@
 from mysql_mimic.errors import MysqlError, ErrorCode
 
 
-class Default:
-    ...
+class Default: ...
 
 
 VariableType = Callable[[Any], Any]
@@ -102,8 +101,7 @@ def list(self) -> list[tuple[str, str]]:
 
     @property
     @abc.abstractmethod
-    def schema(self) -> dict[str, VariableSchema]:
-        ...
+    def schema(self) -> dict[str, VariableSchema]: ...
 
 
 class GlobalVariables(Variables):

tests/conftest.py

@@ -122,8 +122,8 @@ async def to_thread(func: Callable[..., T], *args: Any, **kwargs: Any) -> T:
     return await loop.run_in_executor(None, func_call)
 
 
-@pytest.fixture
-def session() -> MockSession:
+@pytest_asyncio.fixture
+async def session() -> MockSession:
     return MockSession()