kiba-core 0.5.3.dev89__tar.gz → 0.5.3.dev90__tar.gz

{kiba_core-0.5.3.dev89 → kiba_core-0.5.3.dev90}/CHANGELOG.md

@@ -8,6 +8,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 ## [Unreleased]
 
 ### Added
+- [PATCH] Fixed `DatabaseConnectionMiddleware`'s commit when the response starts failing ("another operation is in progress") for streaming responses behind a `BaseHTTPMiddleware`: `Database.execute` and `Database.commit_context_connection` now wait for each other on the same connection
 - [MINOR] Added `Database.commit_context_connection` to commit the context connection's writes so far; anything run on it afterwards is committed when the context connection ends
 - [PATCH] Updated `DatabaseConnectionMiddleware` to commit before the response starts, so a client can't get the response and make its next request before the request's writes are visible (writes made while streaming are still committed when the request ends)
 - [PATCH] Fixed `json_route` and `streaming_json_route` leaving the temporary files of multipart uploads open

{kiba_core-0.5.3.dev89 → kiba_core-0.5.3.dev90}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kiba-core
-Version: 0.5.3.dev89
+Version: 0.5.3.dev90
 Summary: Kiba Labs' python utilities
 Project-URL: repository, https://github.com/kibalabs/core-py
 Requires-Python: <4,>=3.12

{kiba_core-0.5.3.dev89 → kiba_core-0.5.3.dev90}/core/store/database.py

@@ -1,3 +1,4 @@
+import asyncio
 import contextlib
 import contextvars
 import datetime
@@ -47,6 +48,7 @@ class Database:
         self._engine: AsyncEngine | None = None
         self._connectionContext = contextvars.ContextVar[DatabaseConnection | None]('_connectionContext')
         self._connectionsWithWrites = weakref.WeakSet[DatabaseConnection]()
+        self._connectionLocks = weakref.WeakKeyDictionary[DatabaseConnection, asyncio.Lock]()
 
     async def connect(self, poolSize: int = 100) -> None:
         if not self._engine:
@@ -86,6 +88,16 @@ class Database:
         async with self._engine.begin() as connection:
             yield connection
 
+    # NOTE(krishan711): a connection can only run one operation at a time, so execute and commit_context_connection wait
+    # for each other. e.g. DatabaseConnectionMiddleware commits when the response starts, which (behind a BaseHTTPMiddleware)
+    # can be while the route, in another task, is still running queries for a streaming response
+    def _get_connection_lock(self, connection: DatabaseConnection) -> asyncio.Lock:
+        lock = self._connectionLocks.get(connection)
+        if lock is None:
+            lock = asyncio.Lock()
+            self._connectionLocks[connection] = lock
+        return lock
+
     def _get_context_connection(self) -> DatabaseConnection | None:
         try:
             connection = self._connectionContext.get()
@@ -128,8 +140,9 @@ class Database:
         connection = self._get_context_connection()
         if not connection:
             raise InternalServerErrorException(message='No connection found. Please call create_context_connection() for the context.')
-        await connection.commit()
-        self._connectionsWithWrites.discard(connection)
+        async with self._get_connection_lock(connection=connection):
+            await connection.commit()
+            self._connectionsWithWrites.discard(connection)
 
     # NOTE(krishan711): unlike create_context_connection this may be used inside an existing context connection.
     # Everything in the block runs on a new transaction that commits when the block exits, independent of the
@@ -164,6 +177,7 @@ class Database:
             connection = self._get_context_connection()
         if not connection:
             raise InternalServerErrorException(message='No connection found. Please provide a connection or call create_context_connection() for the context.')
-        if isinstance(query, UpdateBase):
-            self._connectionsWithWrites.add(connection)
-        return typing.cast(Result[*tuple[typing.Any, ...]], await connection.execute(statement=query))  # type: ignore[explicit-any]
+        async with self._get_connection_lock(connection=connection):
+            if isinstance(query, UpdateBase):
+                self._connectionsWithWrites.add(connection)
+            return typing.cast(Result[*tuple[typing.Any, ...]], await connection.execute(statement=query))  # type: ignore[explicit-any]

{kiba_core-0.5.3.dev89 → kiba_core-0.5.3.dev90}/kiba_core.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kiba-core
-Version: 0.5.3.dev89
+Version: 0.5.3.dev90
 Summary: Kiba Labs' python utilities
 Project-URL: repository, https://github.com/kibalabs/core-py
 Requires-Python: <4,>=3.12

{kiba_core-0.5.3.dev89 → kiba_core-0.5.3.dev90}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "kiba-core"
-version = "0.5.3.dev89"
+version = "0.5.3.dev90"
 description = "Kiba Labs' python utilities"
 readme = "README.md"
 requires-python = ">=3.12, <4"

kiba_core-0.5.3.dev90/tests/api/test_database_connection_middleware.py

@@ -0,0 +1,168 @@
+import asyncio
+import os
+import uuid
+from collections.abc import AsyncIterator
+
+import pytest
+import pytest_asyncio
+import sqlalchemy
+from sqlalchemy import Column
+from sqlalchemy import MetaData
+from sqlalchemy import Table
+from sqlalchemy import insert
+from sqlalchemy import select
+from sqlalchemy import text
+from starlette.applications import Starlette
+from starlette.middleware import Middleware
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.requests import Request
+from starlette.responses import StreamingResponse
+from starlette.routing import Route
+
+from core.api.middleware.database_connection_middleware import DatabaseConnectionMiddleware
+from core.store.database import Database
+
+PSQL_CONNECTION_STRING = os.environ.get('CORE_TEST_PSQL_CONNECTION_STRING')
+
+metadata = MetaData()
+records = Table(f'records_{uuid.uuid4().hex}', metadata, Column('name', sqlalchemy.Text, primary_key=True))
+
+
+@pytest_asyncio.fixture(params=['sqlite', 'postgresql'] if PSQL_CONNECTION_STRING else ['sqlite'])
+async def database(request, tmp_path):
+    connectionString = PSQL_CONNECTION_STRING if request.param == 'postgresql' else Database.create_sqlite_connection_string(str(tmp_path / 'database.sqlite'))
+    database = Database(connectionString=connectionString)
+    await database.connect(poolSize=2)
+    async with database.create_transaction() as connection:
+        await connection.run_sync(metadata.drop_all)
+        await connection.run_sync(metadata.create_all)
+    yield database
+    async with database.create_transaction() as connection:
+        await connection.run_sync(metadata.drop_all)
+    await database.disconnect()
+
+
+async def get_record_names(database: Database) -> list[str]:
+    async with database.create_transaction() as connection:
+        result = await database.execute(select(records.c.name).order_by(records.c.name), connection=connection)
+    return list(result.scalars().all())
+
+
+async def call_app(app, database: Database) -> list[tuple[str, list[str]]]:
+    sentMessages: list[tuple[str, list[str]]] = []
+
+    async def receive():
+        return {'type': 'http.request', 'body': b'', 'more_body': False}
+
+    async def send(message):
+        sentMessages.append((message['type'], await get_record_names(database=database)))
+
+    await DatabaseConnectionMiddleware(app=app, database=database)({'type': 'http', 'method': 'POST', 'path': '/'}, receive, send)
+    return sentMessages
+
+
+@pytest.mark.asyncio
+async def test_writes_are_committed_before_the_response_starts(database):
+    async def app(scope, receive, send):
+        await database.execute(insert(records).values(name='request'))
+        await send({'type': 'http.response.start', 'status': 200, 'headers': []})
+        await send({'type': 'http.response.body', 'body': b'ok'})
+
+    sentMessages = await call_app(app=app, database=database)
+
+    assert sentMessages == [('http.response.start', ['request']), ('http.response.body', ['request'])]
+
+
+@pytest.mark.asyncio
+async def test_writes_while_streaming_are_committed_when_the_request_ends(database):
+    async def app(scope, receive, send):
+        await database.execute(insert(records).values(name='before'))
+        await send({'type': 'http.response.start', 'status': 200, 'headers': []})
+        await database.execute(insert(records).values(name='streaming'))
+        await send({'type': 'http.response.body', 'body': b'ok'})
+
+    sentMessages = await call_app(app=app, database=database)
+
+    assert sentMessages == [('http.response.start', ['before']), ('http.response.body', ['before'])]
+    assert await get_record_names(database=database) == ['before', 'streaming']
+
+
+@pytest.mark.asyncio
+async def test_writes_are_rolled_back_when_the_request_fails_before_responding(database):
+    async def app(scope, receive, send):
+        await database.execute(insert(records).values(name='request'))
+        raise ValueError('failed')
+
+    with pytest.raises(ValueError, match='failed'):
+        await call_app(app=app, database=database)
+
+    assert await get_record_names(database=database) == []
+
+
+async def post_to_app(app) -> tuple[int, bytes]:
+    scope = {
+        'type': 'http', 'asgi': {'version': '3.0'}, 'http_version': '1.1', 'method': 'POST', 'scheme': 'http', 'path': '/', 'raw_path': b'/',
+        'root_path': '', 'query_string': b'', 'headers': [], 'client': ('127.0.0.1', 1), 'server': ('test', 80),
+    }
+    isRequestSent = False
+    responseFinished = asyncio.Event()
+    statusHolder: list[int] = []
+    bodyParts: list[bytes] = []
+
+    async def receive():
+        nonlocal isRequestSent
+        if not isRequestSent:
+            isRequestSent = True
+            return {'type': 'http.request', 'body': b'', 'more_body': False}
+        await responseFinished.wait()
+        return {'type': 'http.disconnect'}
+
+    async def send(message):
+        if message['type'] == 'http.response.start':
+            statusHolder.append(message['status'])
+        elif message['type'] == 'http.response.body':
+            bodyParts.append(message.get('body', b''))
+            if not message.get('more_body', False):
+                responseFinished.set()
+
+    await app(scope, receive, send)
+    responseFinished.set()
+    return (statusHolder or [0])[0], b''.join(bodyParts)
+
+
+@pytest.mark.asyncio
+async def test_streaming_through_base_http_middleware_does_not_query_the_connection_concurrently(database):
+    # NOTE(krishan711): BaseHTTPMiddleware runs the route in another task, so the stream is still running queries when
+    # the response start reaches DatabaseConnectionMiddleware and it commits
+    async def stream() -> AsyncIterator[bytes]:
+        for index in range(20):
+            if database.get_dialect_name() == 'postgresql':
+                await database.execute(text('SELECT pg_sleep(0.005)'))
+            await database.execute(insert(records).values(name=f'event-{index:02d}'))
+            yield b'event\n'
+
+    async def endpoint(request: Request) -> StreamingResponse:
+        await database.execute(insert(records).values(name='request'))
+        return StreamingResponse(stream())
+
+    async def passthrough(request, callNext):
+        return await callNext(request)
+
+    app = Starlette(
+        routes=[Route('/', endpoint=endpoint, methods=['POST'])],
+        middleware=[
+            Middleware(DatabaseConnectionMiddleware, database=database),
+            Middleware(BaseHTTPMiddleware, dispatch=passthrough),
+        ],
+    )
+    for _ in range(5):
+        status, body = await post_to_app(app=app)
+        assert status == 200
+        assert body == b'event\n' * 20
+        assert len(await get_record_names(database=database)) == 21
+        async with database.create_transaction() as connection:
+            await connection.execute(records.delete())
+
+    async with database.create_transaction() as connection:
+        await connection.execute(insert(records).values(name='check'))
+    assert await get_record_names(database=database) == ['check']

kiba_core-0.5.3.dev89/tests/api/test_database_connection_middleware.py

@@ -1,81 +0,0 @@
-import pytest
-import pytest_asyncio
-import sqlalchemy
-from sqlalchemy import Column
-from sqlalchemy import MetaData
-from sqlalchemy import Table
-from sqlalchemy import insert
-from sqlalchemy import select
-
-from core.api.middleware.database_connection_middleware import DatabaseConnectionMiddleware
-from core.store.database import Database
-
-metadata = MetaData()
-records = Table('records', metadata, Column('name', sqlalchemy.Text, primary_key=True))
-
-
-@pytest_asyncio.fixture
-async def database(tmp_path):
-    database = Database(connectionString=Database.create_sqlite_connection_string(str(tmp_path / 'database.sqlite')))
-    await database.connect(poolSize=2)
-    async with database.create_transaction() as connection:
-        await connection.run_sync(metadata.create_all)
-    yield database
-    await database.disconnect()
-
-
-async def get_record_names(database: Database) -> list[str]:
-    async with database.create_transaction() as connection:
-        result = await database.execute(select(records.c.name).order_by(records.c.name), connection=connection)
-    return list(result.scalars().all())
-
-
-async def call_app(app, database: Database) -> list[tuple[str, list[str]]]:
-    sentMessages: list[tuple[str, list[str]]] = []
-
-    async def receive():
-        return {'type': 'http.request', 'body': b'', 'more_body': False}
-
-    async def send(message):
-        sentMessages.append((message['type'], await get_record_names(database=database)))
-
-    await DatabaseConnectionMiddleware(app=app, database=database)({'type': 'http', 'method': 'POST', 'path': '/'}, receive, send)
-    return sentMessages
-
-
-@pytest.mark.asyncio
-async def test_writes_are_committed_before_the_response_starts(database):
-    async def app(scope, receive, send):
-        await database.execute(insert(records).values(name='request'))
-        await send({'type': 'http.response.start', 'status': 200, 'headers': []})
-        await send({'type': 'http.response.body', 'body': b'ok'})
-
-    sentMessages = await call_app(app=app, database=database)
-
-    assert sentMessages == [('http.response.start', ['request']), ('http.response.body', ['request'])]
-
-
-@pytest.mark.asyncio
-async def test_writes_while_streaming_are_committed_when_the_request_ends(database):
-    async def app(scope, receive, send):
-        await database.execute(insert(records).values(name='before'))
-        await send({'type': 'http.response.start', 'status': 200, 'headers': []})
-        await database.execute(insert(records).values(name='streaming'))
-        await send({'type': 'http.response.body', 'body': b'ok'})
-
-    sentMessages = await call_app(app=app, database=database)
-
-    assert sentMessages == [('http.response.start', ['before']), ('http.response.body', ['before'])]
-    assert await get_record_names(database=database) == ['before', 'streaming']
-
-
-@pytest.mark.asyncio
-async def test_writes_are_rolled_back_when_the_request_fails_before_responding(database):
-    async def app(scope, receive, send):
-        await database.execute(insert(records).values(name='request'))
-        raise ValueError('failed')
-
-    with pytest.raises(ValueError, match='failed'):
-        await call_app(app=app, database=database)
-
-    assert await get_record_names(database=database) == []