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

{kiba_core-0.5.3.dev89 → kiba_core-0.5.3.dev91}/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.dev91}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kiba-core
-Version: 0.5.3.dev89
+Version: 0.5.3.dev91
 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.dev91}/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.dev91}/core/store/entity_repository.py

@@ -9,6 +9,7 @@ from sqlalchemy.dialects import sqlite as sqlalchemy_sqlite
 from sqlalchemy.engine import Result
 from sqlalchemy.engine import RowMapping
 from sqlalchemy.sql import Select
+from sqlalchemy.sql.base import ReadOnlyColumnCollection
 
 from core.exceptions import KibaException
 from core.exceptions import NotFoundException
@@ -25,6 +26,7 @@ from core.store.retriever import uuid_from_value
 from core.util import date_util
 
 EntityType = typing.TypeVar('EntityType', bound=BaseModel)
+UpsertUpdateWhere = typing.Callable[[ReadOnlyColumnCollection[str, typing.Any]], sqlalchemy.ColumnElement[bool]]  # type: ignore[explicit-any]
 
 
 class EntityRepository(typing.Generic[EntityType]):  # noqa: UP046
@@ -135,7 +137,8 @@ class EntityRepository(typing.Generic[EntityType]):  # noqa: UP046
         result = await database.execute(query=self.table.update().where(self.idColumn == idValue).values(updateValues).returning(*self.list_select_columns()), connection=connection)
         return self.force_from_result(result=result)
 
-    async def upsert(self, database: Database, constraintColumnNames: list[str], connection: DatabaseConnection | None = None, **kwargs) -> EntityType:  # type: ignore[no-untyped-def]  # noqa: ANN003
+    # NOTE(krishan711): updateWhere gets the conflicting row's new values (`excluded`) and only updates when it is true; when it is false nothing is written and the stored row is returned
+    async def upsert(self, database: Database, constraintColumnNames: list[str], connection: DatabaseConnection | None = None, updateWhere: UpsertUpdateWhere | None = None, **kwargs) -> EntityType:  # type: ignore[no-untyped-def]  # noqa: ANN003
         updateValues = self._create_values(kwargs=kwargs, shouldAddUpdatedDate=True, shouldAddCreatedDate=False)
         insertValues = self._create_values(kwargs=kwargs, shouldAddUpdatedDate=True, shouldAddCreatedDate=True)
         self._add_generated_id(values=insertValues)
@@ -144,9 +147,16 @@ class EntityRepository(typing.Generic[EntityType]):  # noqa: UP046
         doUpdateStatement = insertStatement.on_conflict_do_update(
             index_elements=constraintColumns,
             set_=updateValues,
+            where=updateWhere(insertStatement.excluded) if updateWhere is not None else None,
         )
         result = await database.execute(query=doUpdateStatement.returning(*self.list_select_columns()), connection=connection)
-        return self.force_from_result(result=result)
+        row = result.mappings().first()
+        if row is None and updateWhere is not None:
+            storedQuery = self.build_select().where(*[column == insertValues[column] for column in constraintColumns])
+            row = (await database.execute(query=storedQuery, connection=connection)).mappings().first()
+        if row is None:
+            raise NotFoundException
+        return self.from_row(row=row)
 
     async def upsert_many(self, database: Database, constraintColumnNames: list[str], rowDicts: list[dict[str, typing.Any]], connection: DatabaseConnection | None = None) -> list[EntityType]:  # type: ignore[explicit-any]
         if not rowDicts:

{kiba_core-0.5.3.dev89 → kiba_core-0.5.3.dev91}/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.dev91
 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.dev91}/pyproject.toml

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

kiba_core-0.5.3.dev91/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 → kiba_core-0.5.3.dev91}/tests/test_entity_repository.py

@@ -172,6 +172,49 @@ async def test_upsert_inserts_then_updates(database: Database):
     assert upsertedItem.count == 2
 
 
+async def test_upsert_with_false_update_where_keeps_and_returns_the_stored_row(database: Database):
+    item = await _create_item(database=database, name='a', count=1)
+    async with database.create_transaction() as connection:
+        upsertedItem = await ItemsRepository.upsert(database=database, connection=connection, constraintColumnNames=['name'], updateWhere=lambda excluded: excluded.count > ItemsTable.c.count, name='a', count=0, amount=5)
+        storedItem = await ItemsRepository.get(database=database, connection=connection, idValue=item.itemId)
+    assert upsertedItem == item
+    assert storedItem == item
+
+
+async def test_upsert_with_false_update_where_finds_the_stored_row_by_converted_constraint_values(database: Database):
+    repository = UpperNameEntityRepository(table=ItemsTable, modelClass=Item)
+    async with database.create_transaction() as connection:
+        item = await repository.create(database=database, connection=connection, name='a', count=1, amount=1)
+        upsertedItem = await repository.upsert(database=database, connection=connection, constraintColumnNames=['name'], updateWhere=lambda _excluded: sqlalchemy.false(), name='a', count=2, amount=1)
+    assert upsertedItem == item
+
+
+async def test_upsert_with_true_update_where_updates_the_row(database: Database):
+    item = await _create_item(database=database, name='a', count=1)
+    async with database.create_transaction() as connection:
+        upsertedItem = await ItemsRepository.upsert(database=database, connection=connection, constraintColumnNames=['name'], updateWhere=lambda excluded: excluded.count > ItemsTable.c.count, name='a', count=2, amount=5)
+    assert (upsertedItem.itemId, upsertedItem.createdDate, upsertedItem.count, upsertedItem.amount) == (item.itemId, item.createdDate, 2, 5)
+
+
+async def test_upsert_with_update_where_still_inserts_new_rows(database: Database):
+    async with database.create_transaction() as connection:
+        upsertedItem = await ItemsRepository.upsert(database=database, connection=connection, constraintColumnNames=['name'], updateWhere=lambda _excluded: sqlalchemy.false(), name='a', count=1, amount=1)
+        storedItems = await ItemsRepository.list_many(database=database, connection=connection, fieldFilters=[])
+    assert storedItems == [upsertedItem]
+
+
+async def test_upsert_only_when_values_changed(database: Database):
+    def changed(excluded: typing.Any) -> sqlalchemy.ColumnElement[bool]:  # type: ignore[explicit-any]
+        return sqlalchemy.or_(ItemsTable.c.count.is_distinct_from(excluded.count), ItemsTable.c.expiryDate.is_distinct_from(excluded.expiryDate))
+
+    async with database.create_transaction() as connection:
+        item = await ItemsRepository.upsert(database=database, connection=connection, constraintColumnNames=['name'], updateWhere=changed, name='a', count=1, amount=1, expiryDate=None)
+        unchangedItem = await ItemsRepository.upsert(database=database, connection=connection, constraintColumnNames=['name'], updateWhere=changed, name='a', count=1, amount=1, expiryDate=None)
+        changedItem = await ItemsRepository.upsert(database=database, connection=connection, constraintColumnNames=['name'], updateWhere=changed, name='a', count=1, amount=1, expiryDate=datetime.datetime(2030, 1, 1, tzinfo=datetime.UTC))
+    assert unchangedItem == item
+    assert (changedItem.itemId, changedItem.expiryDate) == (item.itemId, datetime.datetime(2030, 1, 1, tzinfo=datetime.UTC))
+
+
 async def test_upsert_many_inserts_and_updates(database: Database):
     item = await _create_item(database=database, name='a')
     async with database.create_transaction() as connection:

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) == []