scaffold-framework 0.1.0__py3-none-any.whl

scaffold/cli/__init__.py

@@ -0,0 +1,8 @@
+from .base_app import BaseCLIApp
+from .decorators import argument, command
+
+__all__ = [
+    "BaseCLIApp",
+    "argument",
+    "command",
+]

scaffold/cli/base_app.py

@@ -0,0 +1,37 @@
+import argparse
+import asyncio
+
+from .decorators import Command
+
+
+class BaseCLIApp:
+    def __init__(self) -> None:
+        self.parser = argparse.ArgumentParser()
+        self.subparsers = self.parser.add_subparsers()
+        self._register_commands()
+
+    def _register_commands(self) -> None:
+        for attr_name in dir(self):
+            attr = getattr(self, attr_name)
+            if isinstance(attr, Command):
+                subparser = self.subparsers.add_parser(
+                    attr.command_name,
+                    help=attr.command_help,
+                )
+                for args, kwargs in attr.arguments:
+                    subparser.add_argument(*args, **kwargs)
+                subparser.set_defaults(func=attr)
+
+    def run(self) -> None:
+        args = self.parser.parse_args()
+
+        if hasattr(args, "func"):
+            kwargs = dict(vars(args))
+            del kwargs["func"]
+            result = args.func(self, **kwargs)
+
+            if asyncio.iscoroutine(result):
+                asyncio.run(result)
+
+        else:
+            self.parser.print_help()

scaffold/cli/decorators.py

@@ -0,0 +1,47 @@
+from collections.abc import Callable
+from typing import Any
+
+
+class Command:
+    def __init__(
+        self,
+        func: Callable,
+        name: str = "",
+        description: str = "",
+    ) -> None:
+        self.func = func
+        self.command_name = name
+        self.command_help = description
+        self.arguments: list[tuple[tuple[str, ...], dict[str, Any]]] = []
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
+        return self.func(*args, **kwargs)
+
+
+def command(
+    name: str,
+    description: str = "",
+) -> Callable[[Callable | Command], Command]:
+    def decorator(f: Callable | Command) -> Command:
+        if isinstance(f, Command):
+            f.command_name = name
+            f.command_help = description
+            return f
+
+        return Command(f, name, description)
+
+    return decorator
+
+
+def argument(
+    *args: Any,  # noqa: ANN401
+    **kwargs: Any,  # noqa: ANN401
+) -> Callable[[Callable | Command], Command]:
+    def decorator(f: Callable | Command) -> Command:
+        if not isinstance(f, Command):
+            f = Command(f)
+
+        f.arguments.append((args, kwargs))
+        return f
+
+    return decorator

scaffold/di.py

@@ -0,0 +1,81 @@
+"""
+An extremely simple DI container (in <100 lines of code) with auto-wiring based on type hints.
+"""
+
+import inspect
+from collections.abc import Awaitable, Callable
+from typing import Self, cast
+
+# Since `type` does not accept abstract classes, we have to use `Callable` as well (although it's not ideal), see https://github.com/python/mypy/issues/4717
+type Dependency[T] = type[T] | Callable[..., T]
+type Provider[T] = type[T] | Callable[["Container"], T]
+
+
+class Container:
+    def __init__(self) -> None:
+        self.providers: dict[Dependency, Provider] = {}
+        self.singletons: dict[Dependency, object] = {}
+        self.init_functions: list[Callable] = []
+
+    def add_singleton[T](self, cls: Dependency[T], provider: Provider[T]) -> None:
+        self.providers[cls] = provider
+        self.singletons[cls] = None
+
+    def add_transient[T](self, cls: Dependency[T], provider: Provider[T]) -> None:
+        self.providers[cls] = provider
+
+    def add_init_function(
+        self,
+        init_function: Callable[[Self], Awaitable[None]],
+    ) -> None:
+        self.init_functions.append(init_function)
+
+    async def init(self) -> None:
+        for init_function in self.init_functions:
+            await init_function(self)
+
+    def resolve[T](self, cls: Dependency[T]) -> T:
+        instance: T
+
+        if cls in self.providers:
+            provider = self.providers[cls]
+
+            if cls in self.singletons and self.singletons[cls] is not None:
+                return cast(T, self.singletons[cls])
+
+            if isinstance(provider, type):
+                instance = self._instantiate(provider)
+
+            else:
+                instance = provider(self)
+
+            if cls in self.singletons:
+                self.singletons[cls] = instance
+
+            return instance
+
+        if isinstance(cls, type):
+            return self._instantiate(cls)
+
+        raise RuntimeError
+
+    def _instantiate[T](self, cls: type[T]) -> T:
+        constructor_signature = inspect.signature(cls.__init__)
+        dependencies = {}
+
+        for name, param in constructor_signature.parameters.items():
+            # TODO is there a better way to ignore the "self" param?
+            if name == "self" or param.annotation == inspect.Parameter.empty:
+                continue  # Skip parameters that are not type-annotated or are 'self'
+            dependencies[name] = self[param.annotation]
+
+        return cls(**dependencies)
+
+    def __getitem__[T](self, cls: Dependency[T]) -> T:
+        return self.resolve(cls)
+
+    def get_factory[C](self, cls: Dependency[C]) -> Callable[[], C]:
+        def factory() -> C:
+            return self[cls]
+
+        return factory

scaffold/persistence.py

@@ -0,0 +1,109 @@
+import abc
+from collections.abc import Sequence
+from types import TracebackType, get_original_bases
+from typing import Protocol, final, get_args, override
+
+from sqlalchemy.ext.asyncio import AsyncSession
+
+
+class EntityId(Protocol):
+    @property
+    def value(self) -> object: ...
+
+
+class Entity(Protocol):
+    @property
+    def id(self) -> EntityId: ...
+
+
+class GenericSqlRepository[E: Entity, ID: EntityId, DTO](abc.ABC):
+    dto_class: type[DTO]
+
+    @override
+    def __init_subclass__(cls) -> None:
+        # TODO Check that the runtime type of the ID type param is the same as the type hint of E.id.
+        # Ideally, we would like to do something like `GenericSqlRepository[ID: EntityId, E: Entity[ID], DTO](abc.ABC)`
+        # to check it statically but that's currently not possible.
+
+        cls.dto_class = get_args(get_original_bases(cls)[0])[2]
+
+        return super().__init_subclass__()
+
+    def __init__(self, session: AsyncSession) -> None:
+        self._session = session
+        self._identity_map: dict[EntityId, E] = {}
+
+    async def get(self, entity_id: ID) -> E | None:
+        if entity_id in self._identity_map:
+            return self._identity_map[entity_id]
+
+        dto = await self._session.get(self.dto_class, entity_id.value)
+
+        if dto is not None:
+            return self.map_dto_to_entity_and_track(dto)
+
+        return None
+
+    def add(self, entity: E) -> None:
+        dto = self._map_entity_to_dto(entity)
+        self._session.add(dto)
+        self._track(entity)
+
+    async def remove(self, entity: E) -> None:
+        dto = self._map_entity_to_dto(entity)
+        await self._session.delete(dto)
+        self._identity_map.pop(entity.id, None)
+
+    @final
+    def map_dto_to_entity_and_track(self, dto: DTO) -> E:
+        entity = self._map_dto_to_entity(dto)
+        self._track(entity)
+        return entity
+
+    @abc.abstractmethod
+    def _map_entity_to_dto(self, entity: E) -> DTO:
+        """Convert a domain entity to a DTO."""
+        pass
+
+    @abc.abstractmethod
+    def _map_dto_to_entity(self, dto: DTO) -> E:
+        """Convert a DTO to a domain entity."""
+        pass
+
+    def _track(self, entity: E) -> None:
+        self._identity_map[entity.id] = entity
+
+    async def sync_state(self) -> None:
+        for entity in self._identity_map.values():
+            dto = self._map_entity_to_dto(entity)
+            await self._session.merge(dto)
+
+
+class GenericSqlUnitOfWork:
+    def __init__(self, session: AsyncSession) -> None:
+        self._session = session
+
+    async def __aenter__(self) -> None:
+        return
+
+    async def __aexit__(
+        self,
+        exc_type: type,
+        exc: BaseException,
+        tb: TracebackType,
+    ) -> None:
+        await self.rollback()
+        await self._session.close()
+
+    async def commit(self) -> None:
+        for repo in self._repositories:
+            await repo.sync_state()
+
+        await self._session.commit()
+
+    async def rollback(self) -> None:
+        await self._session.rollback()
+
+    @property
+    def _repositories(self) -> Sequence[GenericSqlRepository]:
+        return [value for value in self.__dict__.values() if isinstance(value, GenericSqlRepository)]

scaffold/pub_sub.py

@@ -0,0 +1,88 @@
+import asyncio
+import json
+from collections.abc import AsyncGenerator
+
+from psycopg import sql
+from psycopg_pool import AsyncConnectionPool
+
+
+class PostgresPubSubService:
+    def __init__(
+        self,
+        connection_pool: AsyncConnectionPool,
+        database_channel: str = "pub_sub_messages",
+    ) -> None:
+        self._connection_pool = connection_pool
+        self._database_channel = database_channel
+        self._subscribers: dict[str, list[asyncio.Queue]] = {}
+        self._listener_task: asyncio.Task | None = None
+
+    async def init(self) -> None:
+        await self._connection_pool.open()
+
+    async def publish(self, channel_name: str, message: str) -> None:
+        async with self._connection_pool.connection() as conn:
+            payload = json.dumps({"channel_name": channel_name, "message": message})
+            await conn.execute(
+                sql.SQL("NOTIFY {database_channel}, {payload}").format(
+                    database_channel=sql.Identifier(self._database_channel),
+                    payload=sql.Literal(payload),
+                ),
+            )
+            await conn.commit()
+
+    async def subscribe(self, channel_name: str) -> AsyncGenerator[str, None]:
+        queue: asyncio.Queue = asyncio.Queue()
+        self._subscribers.setdefault(channel_name, []).append(queue)
+
+        # Start the listener task if it's not already running
+        if not self._listener_task or self._listener_task.done():
+            self._listener_task = asyncio.create_task(self._listen())
+
+        try:
+            while True:
+                message = await queue.get()
+                yield message
+
+        finally:
+            # Cleanup when the subscriber is done
+            self._subscribers[channel_name].remove(queue)
+            if not self._subscribers[channel_name]:
+                del self._subscribers[channel_name]
+
+            # Cancel the listener task if no subscribers remain
+            if not any(self._subscribers.values()) and self._listener_task:
+                self._listener_task.cancel()
+                try:
+                    await self._listener_task
+                except asyncio.CancelledError:
+                    pass  # Listener task has been cancelled
+                self._listener_task = None
+
+    async def _listen(self) -> None:
+        async with self._connection_pool.connection() as conn:
+            await conn.execute(
+                sql.SQL("LISTEN {database_channel}").format(
+                    database_channel=sql.Identifier(self._database_channel),
+                ),
+            )
+            await conn.commit()
+
+            async for notification in conn.notifies():
+                payload = notification.payload
+                try:
+                    # Parse the JSON payload
+                    data = json.loads(payload)
+                    channel_name = data.get("channel_name")
+                    message = data.get("message")
+
+                    if channel_name and message:
+                        # Dispatch the message to subscribers of the logical channel
+                        subscribers = self._subscribers.get(channel_name, [])
+                        for queue in subscribers:
+                            await queue.put(message)
+
+                except json.JSONDecodeError:
+                    # Handle invalid JSON payloads
+                    # TODO log this?
+                    continue

scaffold/task_queue.py

@@ -0,0 +1,211 @@
+import asyncio
+import datetime
+import importlib
+import uuid
+from collections.abc import AsyncGenerator, Callable
+from typing import Protocol, override, runtime_checkable
+
+import pydantic
+from psycopg import sql
+from psycopg.rows import dict_row
+from psycopg.types.json import Json
+from psycopg_pool import AsyncConnectionPool
+
+from scaffold.uuid7 import uuid7
+
+
+@runtime_checkable
+class HandlerProtocol[T](Protocol):
+    async def handle_task(self, task: T) -> None: ...
+
+
+class PostgresTaskQueue[T, H: HandlerProtocol]:
+    @override
+    def __init_subclass__(cls) -> None:
+        # TODO Check that the runtime type of the `T` type param is the same as the type hint of `H.handle_task` `task` param.
+        # Ideally, we would like to do something like `GenericFakeTaskQueue[T, H: HandlerProtocol[T]]`
+        # to check it statically but that's currently not possible. Now, the `HandlerProtocol`` is not parametrized so it's
+        # equivalent to `HandlerProtocol[Any]``.
+        pass
+
+    def __init__(
+        self,
+        connection_pool: AsyncConnectionPool,
+        schema_name: str = "public",
+        table_name: str = "task",
+        notify_channel_name: str = "task_queue_notifications",
+    ) -> None:
+        self._handler_factories: dict[type[T], Callable[[], H]] = {}
+        self._connection_pool = connection_pool
+        self._schema_name = schema_name
+        self._table_name = table_name
+        self._notify_channel_name = notify_channel_name
+
+    async def init(self) -> None:
+        await self._connection_pool.open()
+        async with self._connection_pool.connection() as conn:
+            # TODO add queue name
+            stmt = sql.SQL(
+                """
+            CREATE TABLE IF NOT EXISTS {table} (
+                id UUID PRIMARY KEY,
+                class_name VARCHAR NOT NULL,
+                module_name VARCHAR NOT NULL,
+                data JSONB NOT NULL,
+                enqueued_at TIMESTAMP NOT NULL,
+                dequeued_at TIMESTAMP,
+                acknowledged_at TIMESTAMP,
+                visibility_timeout INTEGER NOT NULL
+            )
+            """,
+            ).format(table=self._full_table_identifier)
+            # TODO create table task_failure
+            await conn.execute(stmt)
+
+    async def enqueue(
+        self,
+        task: T,
+        visibility_timeout: int = 30,
+    ) -> None:
+        async with self._connection_pool.connection() as conn:
+            stmt = sql.SQL(
+                """\
+            INSERT INTO {table} (id, class_name, module_name, data, enqueued_at, visibility_timeout)
+            VALUES (%s, %s, %s, %s, %s, %s)
+            """,
+            ).format(table=self._full_table_identifier)
+
+            class_name = task.__class__.__name__
+            module_name = task.__class__.__module__
+
+            # TODO check that the task is a data class
+            # dataclasses.is_dataclass(task)
+
+            # TODO cache type adapters?
+            data = Json(
+                task,
+                dumps=lambda obj: pydantic.TypeAdapter(task.__class__).dump_json(obj),
+            )
+
+            await conn.execute(
+                stmt,
+                (
+                    str(uuid7()),
+                    class_name,
+                    module_name,
+                    data,
+                    datetime.datetime.utcnow(),
+                    visibility_timeout,
+                ),
+            )
+            await conn.execute(
+                sql.SQL("NOTIFY {channel_name}").format(
+                    channel_name=sql.Identifier(self._notify_channel_name),
+                ),
+            )
+            await conn.commit()
+
+    async def handle_task(self, task_id: uuid.UUID, task: T) -> None:
+        # TODO handle exceptions
+        handler = self._handler_factories[type(task)]()
+        # TODO check if whether the handler is a coroutine function
+        await handler.handle_task(task)
+        await self.ack(task_id)
+
+    async def handle_tasks(self) -> None:
+        async with asyncio.TaskGroup() as tg:
+            async for task_data in self._tasks:
+                task_id = task_data["id"]
+                # TODO cache this
+                task_module = importlib.import_module(task_data["module_name"])
+                task_class = getattr(task_module, task_data["class_name"])
+                task = pydantic.TypeAdapter(task_class).validate_python(
+                    task_data["data"],
+                )
+                tg.create_task(self.handle_task(task_id, task))
+
+    @property
+    async def _tasks(self) -> AsyncGenerator[dict]:
+        async with self._connection_pool.connection() as listen_conn:
+            await listen_conn.execute(
+                sql.SQL("LISTEN {channel_name}").format(
+                    channel_name=sql.Identifier(self._notify_channel_name),
+                ),
+            )
+            await listen_conn.commit()
+
+            while True:
+                task = await self._get_task()
+                if task:
+                    yield task
+                else:
+                    # We are doing long polling as well because there's no notification when a message times out
+                    async for _ in listen_conn.notifies(timeout=1):
+                        break
+
+    async def _get_task(self) -> dict | None:
+        async with self._connection_pool.connection() as conn:
+            cursor = conn.cursor(row_factory=dict_row)
+            cursor = await cursor.execute(
+                sql.SQL(
+                    """\
+                UPDATE {table}
+                SET dequeued_at = NOW()
+                WHERE id = (
+                    SELECT
+                        id
+                    FROM
+                        {table}
+                    WHERE
+                        acknowledged_at IS NULL
+                        AND (dequeued_at IS NULL OR dequeued_at < NOW() - make_interval(secs => visibility_timeout))
+                    ORDER BY
+                        enqueued_at
+                    FOR UPDATE SKIP LOCKED
+                    LIMIT
+                        1
+                )
+                RETURNING id, class_name, module_name, data
+                """,
+                ).format(table=self._full_table_identifier),
+            )
+
+            return await cursor.fetchone()
+
+    async def ack(self, task_id: uuid.UUID) -> None:
+        async with self._connection_pool.connection() as conn:
+            await conn.execute(
+                sql.SQL(
+                    "UPDATE {table} SET acknowledged_at = NOW() WHERE id = %s",
+                ).format(
+                    table=self._full_table_identifier,
+                ),
+                (task_id,),
+            )
+            await conn.commit()
+
+    def register(self, task_type: type[T], handler_factory: Callable[[], H]) -> None:
+        self._handler_factories[task_type] = handler_factory
+
+    @property
+    def _full_table_identifier(self) -> sql.Identifier:
+        return sql.Identifier(self._schema_name, self._table_name)
+
+
+class GenericFakeTaskQueue[T, H: HandlerProtocol]:
+    def __init__(
+        self,
+    ) -> None:
+        self.handler_factories: dict[type[T], Callable[[], H]] = {}
+        self.queue: list[T] = []
+
+    def enqueue(self, task: T) -> None:
+        self.queue.append(task)
+
+    def register(self, task_type: type[T], handler_factory: Callable[[], H]) -> None:
+        self.handler_factories[task_type] = handler_factory
+
+    async def run(self) -> None:
+        for task in self.queue:
+            handler = self.handler_factories[type(task)]()
+            await handler.handle_task(task)

scaffold/utils.py

@@ -0,0 +1,66 @@
+import inspect
+import os
+import pkgutil
+from importlib import import_module
+from types import ModuleType
+from typing import cast
+from urllib.parse import urlparse, urlunparse
+
+
+def find_subclasses[T: type](package: ModuleType, base_class: T) -> list[T]:
+    subclasses: list[T] = []
+
+    for _, subpath, ispkg in pkgutil.iter_modules(package.__path__):
+        fullpath = package.__name__ + "." + subpath
+        module = import_module(fullpath)
+        if ispkg:
+            subclasses += find_subclasses(module, base_class)
+        for _, obj in inspect.getmembers(module):
+            if (
+                inspect.isclass(obj)
+                and issubclass(obj, base_class)
+                and obj != base_class
+                and obj.__module__ == module.__name__
+            ):
+                subclasses.append(cast(T, obj))
+
+    return subclasses
+
+
+def get_env_flag(variable_name: str) -> bool:
+    val = os.environ.get(variable_name)
+    return bool(val and val.lower() not in {"0", "false", "no"})
+
+
+def get_postgresql_url(for_sqlalchemy: bool = False) -> str:
+    # Get the base URL from the environment variable
+    base_url = os.environ.get("DATABASE_URL")
+
+    if not base_url:
+        message = "DATABASE_URL environment variable is not set"
+        raise ValueError(message)
+
+    # Parse the URL
+    parsed = urlparse(base_url)
+
+    # Assert that the scheme is postgresql
+    assert parsed.scheme == "postgresql", "DATABASE_URL must use the postgresql:// scheme"
+
+    if for_sqlalchemy:
+        # Append the driver name for SQLAlchemy
+        new_scheme = "postgresql+psycopg"
+    else:
+        # Keep the original postgresql scheme
+        new_scheme = "postgresql"
+
+    # Reconstruct the URL with the new scheme
+    return urlunparse(
+        (
+            new_scheme,
+            parsed.netloc,
+            parsed.path,
+            parsed.params,
+            parsed.query,
+            parsed.fragment,
+        ),
+    )