kiba-core 0.5.3.dev63__tar.gz → 0.5.3.dev65__tar.gz

{kiba_core-0.5.3.dev63 → kiba_core-0.5.3.dev65}/.github/workflows/pull-request.yml

@@ -45,9 +45,16 @@ jobs:
           json-file-path: ./type-check-results.json
           check-name: core-type-check
           fail-on-error: false
+      - name: Start postgres
+        run: |
+          docker network create tests
+          docker run --detach --name postgres --network tests --env POSTGRES_PASSWORD=test --env POSTGRES_DB=core postgres:16-alpine
+          # NOTE(krishan711): check over TCP: during initdb the image runs a temporary server on the unix socket only, which reports ready and then shuts down
+          for i in $(seq 1 60); do docker exec postgres pg_isready --host 127.0.0.1 --username postgres --quiet && break; sleep 1; done
+          docker exec postgres pg_isready --host 127.0.0.1 --username postgres
       - name: Run tests
         run: |
-          docker run --name tests $DOCKER_IMAGE make test-ci
+          docker run --name tests --network tests --env CORE_TEST_PSQL_CONNECTION_STRING=postgresql+asyncpg://postgres:test@postgres:5432/core $DOCKER_IMAGE make test-ci
           docker cp tests:/app/test-check-results.json ./test-check-results.json
       - name: Upload test results
         uses: kibalabs/github-action-create-annotations@main

{kiba_core-0.5.3.dev63 → kiba_core-0.5.3.dev65}/CHANGELOG.md

@@ -8,6 +8,11 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 ## [Unreleased]
 
 ### Added
+- [MINOR] Added `Lock` with a lease-based `SqlLock` implementation (`tbl_locks`) and a `with_lock` context manager that keeps the lease alive; acquiring raises `LockedException` (naming the current holder) if not acquired within `maxWaitSeconds`, and `get_owner` reports who holds a lock
+- [MINOR] Added `process_util.get_process_name` to identify the current process (`<containerName>:<containerId>` in docker, `<hostname>:<pid>` otherwise)
+- [MINOR] Added `Database.create_isolated_context_connection` to run a block in its own committed transaction inside an existing context connection; logs `ISOLATED_CONNECTION_AFTER_UNCOMMITTED_WRITES` when opened while the outer context connection has uncommitted writes
+- [MAJOR] Added abstract `MessageQueue.get_message_count` (implemented for SQS, AQS, SQL and Cosmos queues) for the number of messages waiting to be processed
+- [MAJOR] Added abstract `MessageQueue.get_inflight_message_count` (implemented for SQS, SQL and Cosmos queues; AQS raises `NotImplementedError` as Azure does not expose it) for the number of messages currently being processed
 - [MAJOR] Reworked `route` to compose JSON or streaming transport, application-resolved authorization, and rate limiting
 - [MAJOR] Replaced authorization-decorator sequences with named `RouteAuthResolver` policies
 - [MINOR] Added `authJwt`, `authBasic`, and `originIp` to `KibaApiRequest`

{kiba_core-0.5.3.dev63 → kiba_core-0.5.3.dev65}/PKG-INFO

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

kiba_core-0.5.3.dev65/core/locks/lock.py

@@ -0,0 +1,61 @@
+import abc
+import asyncio
+import contextlib
+from abc import ABC
+from collections.abc import AsyncIterator
+
+from core import logging
+from core.locks.model import Lease
+
+
+class Lock(ABC):
+    @abc.abstractmethod
+    async def connect(self) -> None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def disconnect(self) -> None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def acquire(self, name: str, ttlSeconds: int = 60, maxWaitSeconds: float = 0) -> Lease:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def extend(self, lease: Lease, ttlSeconds: int = 60) -> bool:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def release(self, lease: Lease) -> None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def get_owner(self, name: str) -> str | None:
+        raise NotImplementedError
+
+    async def _keep_alive(self, lease: Lease, ttlSeconds: int) -> None:
+        while True:
+            await asyncio.sleep(ttlSeconds / 3)
+            try:
+                isExtended = await self.extend(lease=lease, ttlSeconds=ttlSeconds)
+            except Exception as exception:  # noqa: BLE001
+                # NOTE(krishan711): a transient failure must not kill the keep-alive; the lease simply lapses at expiryDate if it keeps failing
+                logging.error(f'Failed to extend lease {lease.name}:')
+                logging.exception(exception)
+                continue
+            if not isExtended:
+                lease.isLost = True
+                logging.error(f'Lost lease {lease.name}')
+                return
+
+    @contextlib.asynccontextmanager
+    async def with_lock(self, name: str, ttlSeconds: int = 60, maxWaitSeconds: float = 0) -> AsyncIterator[Lease]:
+        lease = await self.acquire(name=name, ttlSeconds=ttlSeconds, maxWaitSeconds=maxWaitSeconds)
+        keepAliveTask = asyncio.create_task(self._keep_alive(lease=lease, ttlSeconds=ttlSeconds))
+        try:
+            yield lease
+        finally:
+            keepAliveTask.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await keepAliveTask
+            await self.release(lease=lease)

kiba_core-0.5.3.dev65/core/locks/model.py

@@ -0,0 +1,15 @@
+import datetime
+
+from pydantic import BaseModel
+
+from core.util import date_util
+
+
+class Lease(BaseModel):
+    name: str
+    token: str
+    expiryDate: datetime.datetime
+    isLost: bool = False
+
+    def is_valid(self) -> bool:
+        return not self.isLost and date_util.datetime_from_now() < self.expiryDate

kiba_core-0.5.3.dev65/core/locks/sql.py

@@ -0,0 +1,96 @@
+import asyncio
+import time
+import uuid
+
+import sqlalchemy
+from sqlalchemy.dialects import postgresql as sqlalchemy_psql
+from sqlalchemy.dialects import sqlite as sqlalchemy_sqlite
+
+from core.exceptions import InternalServerErrorException
+from core.exceptions import LockedException
+from core.locks.lock import Lock
+from core.locks.model import Lease
+from core.store.database import Database
+from core.util import date_util
+from core.util import process_util
+
+# NOTE(krishan711): this table is owned here (not in the consuming app's schema.py) so it can ship with
+# core-py. To include it in an app's own alembic-tracked metadata (so autogenerate creates the table),
+# call `LocksTable.to_metadata(appMetadata)` from the app's schema module.
+LocksMetadata = sqlalchemy.MetaData()
+
+LocksTable = sqlalchemy.Table(
+    'tbl_locks',
+    LocksMetadata,
+    sqlalchemy.Column(key='name', name='name', type_=sqlalchemy.Text, primary_key=True, nullable=False),
+    sqlalchemy.Column(key='lockToken', name='lock_token', type_=sqlalchemy.Text, nullable=False),
+    sqlalchemy.Column(key='acquiredDate', name='acquired_date', type_=sqlalchemy.DateTime(timezone=True), nullable=False),
+    sqlalchemy.Column(key='expiryDate', name='expiry_date', type_=sqlalchemy.DateTime(timezone=True), nullable=False),
+    sqlalchemy.Column(key='owner', name='owner', type_=sqlalchemy.Text, nullable=False),
+)
+
+
+class SqlLock(Lock):
+    def __init__(self, database: Database, table: sqlalchemy.Table = LocksTable, pollIntervalSeconds: float = 1.0, owner: str | None = None) -> None:
+        self.database = database
+        self.table = table
+        self.pollIntervalSeconds = pollIntervalSeconds
+        self.owner = owner or process_util.get_process_name()
+
+    async def connect(self) -> None:
+        pass
+
+    async def disconnect(self) -> None:
+        pass
+
+    async def acquire(self, name: str, ttlSeconds: int = 60, maxWaitSeconds: float = 0) -> Lease:
+        deadline = time.monotonic() + maxWaitSeconds
+        while True:
+            lease = await self._try_acquire(name=name, ttlSeconds=ttlSeconds)
+            if lease is not None:
+                return lease
+            if time.monotonic() >= deadline:
+                owner = await self.get_owner(name=name)
+                raise LockedException(message=f'LOCK_HELD: {name} by {owner}')
+            await asyncio.sleep(min(self.pollIntervalSeconds, max(deadline - time.monotonic(), 0)))
+
+    async def get_owner(self, name: str) -> str | None:
+        async with self.database.create_transaction() as connection:
+            query = sqlalchemy.select(self.table.c.owner).where(self.table.c.name == name).where(self.table.c.expiryDate > date_util.datetime_from_now())
+            result = await self.database.execute(query=query, connection=connection)
+            return result.scalar_one_or_none()
+
+    def _create_insert(self, dialectName: str) -> sqlalchemy_psql.Insert | sqlalchemy_sqlite.Insert:
+        if dialectName == 'postgresql':
+            return sqlalchemy_psql.insert(self.table)
+        if dialectName == 'sqlite':
+            return sqlalchemy_sqlite.insert(self.table)
+        raise InternalServerErrorException(message=f'SqlLock does not support dialect: {dialectName}')
+
+    async def _try_acquire(self, name: str, ttlSeconds: int) -> Lease | None:
+        now = date_util.datetime_from_now()
+        lease = Lease(name=name, token=str(uuid.uuid4()), expiryDate=date_util.datetime_from_now(seconds=ttlSeconds))
+        values = {self.table.c.lockToken: lease.token, self.table.c.acquiredDate: now, self.table.c.expiryDate: lease.expiryDate, self.table.c.owner: self.owner}
+        async with self.database.create_transaction() as connection:
+            upsertQuery = (
+                self._create_insert(dialectName=connection.dialect.name).values({self.table.c.name: name, **values}).on_conflict_do_update(index_elements=[self.table.c.name], set_=values, where=self.table.c.expiryDate <= now).returning(self.table.c.name)
+            )
+            result = await self.database.execute(query=upsertQuery, connection=connection)
+            isAcquired = result.first() is not None
+        return lease if isAcquired else None
+
+    async def extend(self, lease: Lease, ttlSeconds: int = 60) -> bool:
+        now = date_util.datetime_from_now()
+        newExpiryDate = date_util.datetime_from_now(seconds=ttlSeconds)
+        async with self.database.create_transaction() as connection:
+            updateQuery = sqlalchemy.update(self.table).where(self.table.c.name == lease.name).where(self.table.c.lockToken == lease.token).where(self.table.c.expiryDate > now).values(expiryDate=newExpiryDate).returning(self.table.c.name)
+            result = await self.database.execute(query=updateQuery, connection=connection)
+            isExtended = result.first() is not None
+        if isExtended:
+            lease.expiryDate = newExpiryDate
+        return isExtended
+
+    async def release(self, lease: Lease) -> None:
+        async with self.database.create_transaction() as connection:
+            deleteQuery = sqlalchemy.delete(self.table).where(self.table.c.name == lease.name).where(self.table.c.lockToken == lease.token)
+            await self.database.execute(query=deleteQuery, connection=connection)

{kiba_core-0.5.3.dev63 → kiba_core-0.5.3.dev65}/core/queues/aqs.py

@@ -80,3 +80,12 @@ class AqsMessageQueue(MessageQueue[AqsMessage]):
         if not self._aqsClient:
             raise InternalServerErrorException('You need to call .connect() before trying to delete messages')
         await self._aqsClient.delete_message(message=message.aqsId, pop_receipt=message.popReceipt)
+
+    async def get_message_count(self) -> int:
+        if not self._aqsClient:
+            raise InternalServerErrorException('You need to call .connect() before trying to count messages')
+        properties = await self._aqsClient.get_queue_properties()
+        return int(properties.approximate_message_count or 0)
+
+    async def get_inflight_message_count(self) -> int:
+        raise NotImplementedError('Azure Storage Queues do not expose an in-flight message count')

{kiba_core-0.5.3.dev63 → kiba_core-0.5.3.dev65}/core/queues/cosmos.py

@@ -188,3 +188,27 @@ class CosmosMessageQueue(MessageQueue[CosmosMessage]):
             ],
             partition_key=self.queueName,
         )
+
+    async def get_message_count(self) -> int:
+        counts = self.container.query_items(
+            query='SELECT COUNT(1) AS messageCount FROM c WHERE c.itemType = @itemType AND c.queueName = @queueName AND c.visibleDate <= @now',
+            parameters=[
+                {'name': '@itemType', 'value': 'message'},
+                {'name': '@queueName', 'value': self.queueName},
+                {'name': '@now', 'value': time.time()},
+            ],
+            partition_key=self.queueName,
+        )
+        return sum([int(item['messageCount']) async for item in counts])
+
+    async def get_inflight_message_count(self) -> int:
+        counts = self.container.query_items(
+            query='SELECT COUNT(1) AS messageCount FROM c WHERE c.itemType = @itemType AND c.queueName = @queueName AND c.leaseId != null AND c.visibleDate > @now',
+            parameters=[
+                {'name': '@itemType', 'value': 'message'},
+                {'name': '@queueName', 'value': self.queueName},
+                {'name': '@now', 'value': time.time()},
+            ],
+            partition_key=self.queueName,
+        )
+        return sum([int(item['messageCount']) async for item in counts])

{kiba_core-0.5.3.dev63 → kiba_core-0.5.3.dev65}/core/queues/message_queue.py

@@ -33,3 +33,11 @@ class MessageQueue[MessageType: Message](ABC):
     @abc.abstractmethod
     async def delete_message(self, message: MessageType) -> None:
         raise NotImplementedError
+
+    @abc.abstractmethod
+    async def get_message_count(self) -> int:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def get_inflight_message_count(self) -> int:
+        raise NotImplementedError

{kiba_core-0.5.3.dev63 → kiba_core-0.5.3.dev65}/core/queues/sql.py

@@ -16,9 +16,9 @@ from core.queues.model import Message
 from core.store.database import Database
 from core.util import date_util
 
-# NOTE: this table is owned here (not in the consuming app's schema.py) so it can ship with
-# core-py. To include it in an app's own alembic-tracked metadata (so autogenerate creates the
-# table), call `QueueMessagesTable.to_metadata(appMetadata)` from the app's schema module.
+# NOTE(krishan711): this table is owned here (not in the consuming app's schema.py) so it can ship with
+# core-py. To include it in an app's own alembic-tracked metadata (so autogenerate creates the table),
+# call `QueueMessagesTable.to_metadata(appMetadata)` from the app's schema module.
 QueueMessagesMetadata = sqlalchemy.MetaData()
 
 QueueMessagesTable = sqlalchemy.Table(
@@ -163,3 +163,13 @@ class SqlMessageQueue(MessageQueue[SqlMessage]):
         async with self.database.create_transaction() as connection:
             deleteQuery = sqlalchemy.delete(self.table).where(self.table.c.id == message.id).where(self.table.c.lockToken == message.lockToken)
             await self.database.execute(query=deleteQuery, connection=connection)  # type: ignore[arg-type]
+
+    async def get_message_count(self) -> int:
+        countQuery = sqlalchemy.select(sqlalchemy.func.count()).select_from(self.table).where(self.table.c.queueName == self.queueName).where(self.table.c.visibleDate <= date_util.datetime_from_now())
+        result = await self.database.execute(query=countQuery)
+        return int(result.scalar_one())
+
+    async def get_inflight_message_count(self) -> int:
+        countQuery = sqlalchemy.select(sqlalchemy.func.count()).select_from(self.table).where(self.table.c.queueName == self.queueName).where(self.table.c.lockToken.is_not(None)).where(self.table.c.visibleDate > date_util.datetime_from_now())
+        result = await self.database.execute(query=countQuery)
+        return int(result.scalar_one())

{kiba_core-0.5.3.dev63 → kiba_core-0.5.3.dev65}/core/queues/sqs.py

@@ -96,3 +96,15 @@ class SqsMessageQueue(MessageQueue[SqsMessage]):
         if not self._sqsClient:
             raise InternalServerErrorException('You need to call .connect() before trying to delete messages')
         await self._sqsClient.delete_message(QueueUrl=self.queueUrl, ReceiptHandle=message.receiptHandle)
+
+    async def get_message_count(self) -> int:
+        if not self._sqsClient:
+            raise InternalServerErrorException('You need to call .connect() before trying to count messages')
+        response = await self._sqsClient.get_queue_attributes(QueueUrl=self.queueUrl, AttributeNames=['ApproximateNumberOfMessages'])
+        return int(response['Attributes']['ApproximateNumberOfMessages'])
+
+    async def get_inflight_message_count(self) -> int:
+        if not self._sqsClient:
+            raise InternalServerErrorException('You need to call .connect() before trying to count messages')
+        response = await self._sqsClient.get_queue_attributes(QueueUrl=self.queueUrl, AttributeNames=['ApproximateNumberOfMessagesNotVisible'])
+        return int(response['Attributes']['ApproximateNumberOfMessagesNotVisible'])

{kiba_core-0.5.3.dev63 → kiba_core-0.5.3.dev65}/core/store/database.py

@@ -1,6 +1,7 @@
 import contextlib
 import contextvars
 import typing
+import weakref
 from collections.abc import AsyncIterator
 from typing import TypeVar
 
@@ -10,6 +11,7 @@ from sqlalchemy.engine import Result
 from sqlalchemy.ext.asyncio import AsyncConnection
 from sqlalchemy.ext.asyncio import AsyncEngine
 from sqlalchemy.ext.asyncio import create_async_engine
+from sqlalchemy.sql.dml import UpdateBase
 from sqlalchemy.sql.selectable import TypedReturnsRows
 
 from core import logging
@@ -37,6 +39,7 @@ class Database:
         self.connectionString = connectionString
         self._engine: AsyncEngine | None = None
         self._connectionContext = contextvars.ContextVar[DatabaseConnection | None]('_connectionContext')
+        self._connectionsWithWrites = weakref.WeakSet[DatabaseConnection]()
 
     async def connect(self, poolSize: int = 100) -> None:
         if not self._engine:
@@ -99,6 +102,25 @@ class Database:
             await self.disconnect()
             await self.connect()
 
+    # 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
+    # transaction's uncommitted writes and will block forever on rows the outer transaction has written, so
+    # callers must commit their writes before entering. Violations are logged (with the caller's stack) for now.
+    @contextlib.asynccontextmanager
+    async def create_isolated_context_connection(self) -> AsyncIterator[DatabaseConnection]:
+        if not self._engine:
+            raise InternalServerErrorException(message='Engine has not been established. Please called collect() first.')
+        outerConnection = self._get_context_connection()
+        if outerConnection is not None and outerConnection in self._connectionsWithWrites:
+            logging.error('ISOLATED_CONNECTION_AFTER_UNCOMMITTED_WRITES: an isolated context connection was opened while the outer context connection has uncommitted writes', stack_info=True)
+        async with self._engine.begin() as connection:
+            token = self._connectionContext.set(connection)
+            try:
+                yield connection
+            finally:
+                self._connectionContext.reset(token)
+
     @typing.overload
     async def execute(self, query: TypedReturnsRows[ResultType], connection: DatabaseConnection | None = None) -> Result[ResultType]: ...
     @typing.overload
@@ -110,4 +132,6 @@ 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[typing.Any], await connection.execute(statement=query))  # type: ignore[explicit-any]

kiba_core-0.5.3.dev65/core/util/process_util.py

@@ -0,0 +1,28 @@
+import os
+import re
+import socket
+
+DOCKERENV_PATH = '/.dockerenv'
+MOUNTINFO_PATH = '/proc/self/mountinfo'
+_DOCKER_CONTAINER_ID_PATTERN = re.compile(r'/containers/([0-9a-f]{64})/')
+
+
+def is_running_in_docker() -> bool:
+    return os.path.exists(DOCKERENV_PATH)
+
+
+def get_docker_container_id() -> str:
+    try:
+        with open(MOUNTINFO_PATH, encoding='utf-8') as mountInfoFile:
+            match = _DOCKER_CONTAINER_ID_PATTERN.search(mountInfoFile.read())
+    except OSError:
+        match = None
+    # NOTE(krishan711): docker sets the hostname to the short container id unless --hostname is passed
+    return match.group(1)[:12] if match else socket.gethostname()
+
+
+def get_process_name(containerName: str | None = None) -> str:
+    if is_running_in_docker():
+        containerId = get_docker_container_id()
+        return f'{containerName}:{containerId}' if containerName else containerId
+    return f'{socket.gethostname()}:{os.getpid()}'

{kiba_core-0.5.3.dev63 → kiba_core-0.5.3.dev65}/kiba_core.egg-info/PKG-INFO

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

{kiba_core-0.5.3.dev63 → kiba_core-0.5.3.dev65}/kiba_core.egg-info/SOURCES.txt

@@ -44,6 +44,10 @@ core/http/__init__.py
 core/http/basic_authentication.py
 core/http/jwt.py
 core/http/rest_method.py
+core/locks/__init__.py
+core/locks/lock.py
+core/locks/model.py
+core/locks/sql.py
 core/notifications/__init__.py
 core/notifications/discord_client.py
 core/notifications/notification_client.py
@@ -72,6 +76,7 @@ core/util/hashing_util.py
 core/util/http_util.py
 core/util/json_util.py
 core/util/list_util.py
+core/util/process_util.py
 core/util/string_util.py
 core/util/typing_util.py
 core/util/url_util.py
@@ -99,6 +104,7 @@ tests/api/test_route.py
 tests/api/test_streaming_json_route.py
 tests/caching/test_dict_cache.py
 tests/caching/test_file_cache.py
+tests/locks/test_sql_lock.py
 tests/requester/__init__.py
 tests/requester/test_requester.py
 tests/util/__init__.py
@@ -109,6 +115,7 @@ tests/util/test_dict_util.py
 tests/util/test_file_util.py
 tests/util/test_json_util.py
 tests/util/test_list_util.py
+tests/util/test_process_util.py
 tests/util/test_string_util.py
 tests/util/test_url_util.py
 tests/web3/__init__.py

{kiba_core-0.5.3.dev63 → kiba_core-0.5.3.dev65}/pyproject.toml

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