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

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

@@ -8,6 +8,10 @@ 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
 - [MINOR] Updated `json_route` and `streaming_json_route` (and so `route`) to accept `multipart/form-data` bodies: form fields are used like JSON body fields and uploaded files are passed as `core.api.uploaded_file.UploadedFile` (`filename`, `contentType`, `content`), so request models can declare `file: UploadedFile`. A field given more than once raises `BadRequestException`. The `core-api` extra now installs `python-multipart`
 - [MAJOR] Added abstract `MessageQueue.release_message` to hand a claimed message back to be claimed again straight away without using an attempt (`SqlMessageQueue` and `CosmosMessageQueue` keep `postCount`; `SqsMessageQueue` and `AqsMessageQueue` make the message visible again, though SQS still counts the next receive towards its redrive policy)
 - [MAJOR] Added abstract `MessageQueue.delete_expired_messages`: `SqlMessageQueue` deletes `succeeded` and `deduplicated` messages older than `retentionSeconds` (the other queues remove messages on completion and return 0). `SqlMessageQueue.get_message_count` no longer deletes them and its `shouldSkipRemovingNonRetained` parameter is removed

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kiba-core
-Version: 0.5.3.dev88
+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.dev88 → kiba_core-0.5.3.dev90}/core/api/middleware/database_connection_middleware.py

@@ -1,4 +1,5 @@
 from starlette.types import ASGIApp
+from starlette.types import Message
 from starlette.types import Receive
 from starlette.types import Scope
 from starlette.types import Send
@@ -18,4 +19,13 @@ class DatabaseConnectionMiddleware:
             await self.app(scope, receive, send)
             return
         async with self.database.create_context_connection():
-            await self.app(scope, receive, send)
+
+            # NOTE(krishan711): commit before the response starts, otherwise a client can get the response and make its
+            # next request before this request's writes are visible. Anything after this (e.g. while streaming) is
+            # committed when the context connection ends
+            async def send_after_commit(message: Message) -> None:
+                if message['type'] == 'http.response.start':
+                    await self.database.commit_context_connection()
+                await send(message)
+
+            await self.app(scope, receive, send_after_commit)

{kiba_core-0.5.3.dev88 → kiba_core-0.5.3.dev90}/core/api/request_body.py

@@ -12,12 +12,13 @@ from core.util.typing_util import JsonObject
 
 async def read_request_body(request: Request) -> Mapping[str, object]:
     if request.headers.get('Content-Type', '').startswith('multipart/form-data'):
-        form = await request.form()
         body: dict[str, object] = {}
-        for key, value in form.multi_items():
-            if key in body:
-                raise BadRequestException(f'Invalid form body: {key} is given more than once')
-            body[key] = UploadedFile(filename=value.filename, contentType=value.content_type, content=await value.read()) if isinstance(value, UploadFile) else value
+        # NOTE(krishan711): the form is closed on exit so uploaded files' temporary files are closed (their content has been read by then)
+        async with request.form() as form:
+            for key, value in form.multi_items():
+                if key in body:
+                    raise BadRequestException(f'Invalid form body: {key} is given more than once')
+                body[key] = UploadedFile(filename=value.filename, contentType=value.content_type, content=await value.read()) if isinstance(value, UploadFile) else value
         return body
     bodyBytes = await request.body()
     if len(bodyBytes) == 0:

{kiba_core-0.5.3.dev88 → 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()
@@ -107,12 +119,12 @@ class Database:
             raise InternalServerErrorException(message='Engine has not been established. Please called collect() first.')
         if self._get_context_connection() is not None:
             raise InternalServerErrorException(message='Connection has already been established in this context.')
-        connection = None
         try:
-            async with self._engine.begin() as connection:
+            async with self._engine.connect() as connection:
                 self._connectionContext.set(connection)
                 try:
                     yield connection
+                    await self.commit_context_connection()
                 finally:
                     self._connectionContext.set(None)
         except sqlalchemy.exc.InterfaceError as exception:
@@ -122,6 +134,16 @@ class Database:
             await self.disconnect()
             await self.connect()
 
+    # NOTE(krishan711): commits what the context connection has done so far. Anything run on it afterwards is in a new
+    # transaction, which is committed when the context connection ends (or rolled back if it ends with an exception).
+    async def commit_context_connection(self) -> None:
+        connection = self._get_context_connection()
+        if not connection:
+            raise InternalServerErrorException(message='No connection found. Please call create_context_connection() for the context.')
+        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
     # outer transaction, which becomes the context connection again afterwards. It cannot see the outer
@@ -155,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.dev88 → 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.dev88
+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.dev88 → kiba_core-0.5.3.dev90}/kiba_core.egg-info/SOURCES.txt

@@ -106,6 +106,7 @@ tests/test_exceptions.py
 tests/test_logging.py
 tests/api/test_api_response.py
 tests/api/test_authorizer.py
+tests/api/test_database_connection_middleware.py
 tests/api/test_exception_handling_middleware.py
 tests/api/test_json_route.py
 tests/api/test_openapi.py

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

@@ -1,6 +1,6 @@
 [project]
 name = "kiba-core"
-version = "0.5.3.dev88"
+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.dev88 → kiba_core-0.5.3.dev90}/tests/api/test_json_route.py

@@ -1,6 +1,7 @@
 import pytest
 from starlette.applications import Starlette
 from starlette.routing import Route
+from starlette.datastructures import UploadFile
 from starlette.testclient import TestClient
 from pydantic import BaseModel
 
@@ -189,3 +190,21 @@ def test_json_route_with_multipart_body_rejects_repeated_fields(client):
     )
     assert response.status_code == 400
     assert "file is given more than once" in response.json().get("message", "")
+
+
+def test_json_route_with_multipart_body_closes_uploaded_files(client, monkeypatch):
+    closedFilenames = []
+    originalClose = UploadFile.close
+
+    async def close(self):
+        closedFilenames.append(self.filename)
+        await originalClose(self)
+
+    monkeypatch.setattr(UploadFile, "close", close)
+    response = client.post(
+        "/upload",
+        files={"file": ("scan.dcm", b"file-content", "application/dicom")},
+        data={"age": "41"},
+    )
+    assert response.status_code == 200
+    assert closedFilenames == ["scan.dcm"]

{kiba_core-0.5.3.dev88 → kiba_core-0.5.3.dev90}/tests/test_database.py

@@ -229,3 +229,58 @@ async def test_isolated_context_connection_raises_on_uncommitted_outer_writes_wh
             async with database.create_isolated_context_connection():
                 pass
     await database.disconnect()
+
+
+@pytest.mark.asyncio
+async def test_commit_context_connection_makes_writes_visible_before_the_context_ends(database):
+    metadata = MetaData()
+    records = Table('records', metadata, Column('name', sqlalchemy.Text, primary_key=True))
+    await database.connect(poolSize=2)
+    async with database.create_transaction() as connection:
+        await connection.run_sync(metadata.create_all)
+
+    with pytest.raises(ValueError, match='rollback'):
+        async with database.create_context_connection():
+            await database.execute(insert(records).values(name='committed'))
+            await database.commit_context_connection()
+            async with database.create_transaction() as connection:
+                result = await database.execute(select(records.c.name), connection=connection)
+            assert result.scalars().all() == ['committed']
+            await database.execute(insert(records).values(name='rolledback'))
+            raise ValueError('rollback')
+
+    async with database.create_transaction() as connection:
+        result = await database.execute(select(records.c.name), connection=connection)
+    assert result.scalars().all() == ['committed']
+    await database.disconnect()
+
+
+@pytest.mark.asyncio
+async def test_writes_after_commit_context_connection_are_committed_when_the_context_ends(database, caplog):
+    metadata = MetaData()
+    records = Table('records', metadata, Column('name', sqlalchemy.Text, primary_key=True))
+    await database.connect(poolSize=2)
+    async with database.create_transaction() as connection:
+        await connection.run_sync(metadata.create_all)
+
+    with caplog.at_level(logging.ERROR):
+        async with database.create_context_connection():
+            await database.execute(insert(records).values(name='first'))
+            await database.commit_context_connection()
+            async with database.create_isolated_context_connection():
+                pass
+            await database.execute(insert(records).values(name='second'))
+    assert 'ISOLATED_CONNECTION_AFTER_UNCOMMITTED_WRITES' not in caplog.text
+
+    async with database.create_transaction() as connection:
+        result = await database.execute(select(records.c.name).order_by(records.c.name), connection=connection)
+    assert result.scalars().all() == ['first', 'second']
+    await database.disconnect()
+
+
+@pytest.mark.asyncio
+async def test_commit_context_connection_requires_a_context_connection(database):
+    await database.connect(poolSize=1)
+    with pytest.raises(InternalServerErrorException, match='No connection found'):
+        await database.commit_context_connection()
+    await database.disconnect()