Working with async graphql db sessions; Tests broken

Author: Sparrow1029

orchestrator/app.py

@@ -43,7 +43,7 @@
 from orchestrator.api.api_v1.api import api_router
 from orchestrator.api.error_handling import ProblemDetailException
 from orchestrator.cli.main import app as cli_app
-from orchestrator.db import db, init_database
+from orchestrator.db import db, init_async_database, init_database
 from orchestrator.db.database import DBSessionMiddleware
 from orchestrator.db.listeners import monitor_sqlalchemy_queries
 from orchestrator.db.loaders import init_model_loaders
@@ -128,6 +128,7 @@ def __init__(
         self.include_router(api_router, prefix="/api")
 
         init_database(base_settings)
+        init_async_database(base_settings)
 
         self.add_middleware(ClearStructlogContextASGIMiddleware)
         self.add_middleware(SessionMiddleware, secret_key=base_settings.SESSION_SECRET)

orchestrator/db/__init__.py

@@ -14,8 +14,8 @@
 
 from structlog import get_logger
 
+from orchestrator.db.database import AsyncDatabase, Database, transactional
 from orchestrator.db.database import BaseModel as DbBaseModel
-from orchestrator.db.database import Database, transactional
 from orchestrator.db.models import (  # noqa: F401
     EngineSettingsTable,
     FixedInputTable,
@@ -42,15 +42,15 @@
 
 
 class WrappedDatabase:
-    def __init__(self, wrappee: Database | None = None) -> None:
+    def __init__(self, wrappee: Database | AsyncDatabase | None = None) -> None:
         self.wrapped_database = wrappee
 
-    def update(self, wrappee: Database) -> None:
+    def update(self, wrappee: Database | AsyncDatabase) -> None:
         self.wrapped_database = wrappee
         logger.info("Database object configured, all methods referencing `db` should work.")
 
     def __getattr__(self, attr: str) -> Any:
-        if not isinstance(self.wrapped_database, Database):
+        if not isinstance(self.wrapped_database, Database | AsyncDatabase):
             if "_" in attr:
                 logger.warning("No database configured, but attempting to access class methods")
                 return None
@@ -72,6 +72,15 @@ def init_database(settings: AppSettings) -> Database:
     return db
 
 
+async_wrapped_db = WrappedDatabase()
+async_db = cast(AsyncDatabase, async_wrapped_db)
+
+
+def init_async_database(settings: AppSettings) -> AsyncDatabase:
+    async_wrapped_db.update(AsyncDatabase(str(settings.DATABASE_URI)))
+    return async_db
+
+
 __all__ = [
     "transactional",
     "SubscriptionTable",
@@ -94,6 +103,8 @@ def init_database(settings: AppSettings) -> Database:
     "UtcTimestampError",
     "db",
     "init_database",
+    "async_db",
+    "init_async_database",
 ]
 
 ALL_DB_MODELS: list[type[DbBaseModel]] = [

orchestrator/db/database.py

@@ -12,14 +12,16 @@
 # limitations under the License.
 
 from collections.abc import Callable, Generator, Iterator
-from contextlib import contextmanager
+from contextlib import asynccontextmanager, contextmanager
 from contextvars import ContextVar
-from typing import Any, ClassVar, cast
+from typing import Any, AsyncIterator, ClassVar, Mapping, cast
 from uuid import uuid4
 
 import structlog
+from pydantic import PostgresDsn
 from sqlalchemy import create_engine
 from sqlalchemy import inspect as sa_inspect
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
 from sqlalchemy.ext.declarative import DeclarativeMeta
 from sqlalchemy.orm import Query, Session, as_declarative, scoped_session, sessionmaker
 from sqlalchemy.sql.schema import MetaData
@@ -171,6 +173,7 @@ def commit(self) -> None:
     "json_deserializer": json_loads,
 }
 SESSION_ARGUMENTS = {"class_": WrappedSession, "autocommit": False, "autoflush": True, "query_cls": SearchQuery}
+ASYNC_SESSION_ARGUMENTS = {"class_": AsyncSession, "autocommit": False, "autoflush": True, "query_cls": SearchQuery}
 
 
 class Database:
@@ -220,10 +223,23 @@ def database_scope(self, **kwargs: Any) -> Generator["Database", None, None]:
             self.request_context.reset(token)
 
 
+class AsyncDatabase:
+    def __init__(self, db_url: PostgresDsn | str):
+        self.engine = create_async_engine(str(db_url), **ENGINE_ARGUMENTS)
+        self.session_factory = async_sessionmaker(bind=self.engine, expire_on_commit=False, **ASYNC_SESSION_ARGUMENTS)  # type: ignore[call-overload]
+
+    @asynccontextmanager
+    async def session(self, **kwargs: Mapping) -> AsyncIterator[AsyncSession]:
+        async with self.session_factory() as session:
+            yield session
+
+    async def dispose(self) -> None:
+        await self.engine.dispose()
+
+
 class DBSessionMiddleware(BaseHTTPMiddleware):
     def __init__(self, app: ASGIApp, database: Database, commit_on_exit: bool = False):
         super().__init__(app)
-        self.commit_on_exit = commit_on_exit
         self.database = database
 
     async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:

orchestrator/graphql/resolvers/process.py

@@ -17,7 +17,7 @@
 from sqlalchemy import func, select
 from sqlalchemy.orm import selectinload
 
-from orchestrator.db import ProcessTable, db
+from orchestrator.db import ProcessTable
 from orchestrator.db.filters import Filter
 from orchestrator.db.filters.process import PROCESS_TABLE_COLUMN_CLAUSES, filter_processes, process_filter_fields
 from orchestrator.db.models import ProcessSubscriptionTable, SubscriptionTable
@@ -42,7 +42,6 @@
 
 logger = structlog.get_logger(__name__)
 
-
 detailed_props = ("steps", "form", "current_state")
 simple_props = tuple([to_lower_camel(key) for key in ProcessType.__annotations__ if key not in detailed_props])
 
@@ -56,9 +55,10 @@ def _enrich_process(process: ProcessTable, with_details: bool = False) -> Proces
 
 
 async def resolve_process(info: OrchestratorInfo, process_id: UUID) -> ProcessType | None:
+    session = info.context.db_session
     query_loaders = get_query_loaders_for_gql_fields(ProcessTable, info)
     stmt = select(ProcessTable).options(*query_loaders).where(ProcessTable.process_id == process_id)
-    if process := db.session.scalar(stmt):
+    if process := await session.scalar(stmt):
         is_detailed = _is_process_detailed(info)
         return ProcessType.from_pydantic(_enrich_process(process, is_detailed))
     return None
@@ -72,6 +72,7 @@ async def resolve_processes(
     after: int = 0,
     query: str | None = None,
 ) -> Connection[ProcessType]:
+    session = info.context.db_session
     _error_handler = create_resolver_error_handler(info)
     pydantic_filter_by: list[Filter] = [item.to_pydantic() for item in filter_by] if filter_by else []
     pydantic_sort_by: list[Sort] = [item.to_pydantic() for item in sort_by] if sort_by else []
@@ -98,7 +99,7 @@ async def resolve_processes(
         stmt = select_stmt
 
     stmt = sort_processes(stmt, pydantic_sort_by, _error_handler)
-    total = db.session.scalar(select(func.count()).select_from(stmt.subquery()))
+    total = await session.scalar(select(func.count()).select_from(stmt.subquery()))
     stmt = apply_range_to_statement(stmt, after, after + first + 1)
 
     graphql_processes = []

orchestrator/graphql/resolvers/product.py

@@ -1,7 +1,6 @@
 import structlog
 from sqlalchemy import func, select
 
-from orchestrator.db import db
 from orchestrator.db.filters import Filter
 from orchestrator.db.filters.product import PRODUCT_TABLE_COLUMN_CLAUSES, filter_products, product_filter_fields
 from orchestrator.db.models import ProductTable
@@ -28,6 +27,7 @@ async def resolve_products(
     query: str | None = None,
 ) -> Connection[ProductType]:
     _error_handler = create_resolver_error_handler(info)
+    session = info.context.db_session
 
     pydantic_filter_by: list[Filter] = [item.to_pydantic() for item in filter_by] if filter_by else []
     pydantic_sort_by: list[Sort] = [item.to_pydantic() for item in sort_by] if sort_by else []
@@ -49,7 +49,7 @@ async def resolve_products(
         stmt = select_stmt
 
     stmt = sort_products(stmt, pydantic_sort_by, _error_handler)
-    total = db.session.scalar(select(func.count()).select_from(stmt.subquery()))
+    total = await session.scalar(select(func.count()).select_from(stmt.subquery()))
     stmt = apply_range_to_statement(stmt, after, after + first + 1)
 
     graphql_products = []

orchestrator/graphql/resolvers/product_block.py

@@ -1,7 +1,6 @@
 import structlog
 from sqlalchemy import func, select
 
-from orchestrator.db import db
 from orchestrator.db.filters import Filter
 from orchestrator.db.filters.product_block import (
     PRODUCT_BLOCK_TABLE_COLUMN_CLAUSES,
@@ -32,6 +31,7 @@ async def resolve_product_blocks(
     query: str | None = None,
 ) -> Connection[ProductBlock]:
     _error_handler = create_resolver_error_handler(info)
+    session = info.context.db_session
 
     pydantic_filter_by: list[Filter] = [item.to_pydantic() for item in filter_by] if filter_by else []
     pydantic_sort_by: list[Sort] = [item.to_pydantic() for item in sort_by] if sort_by else []
@@ -55,7 +55,7 @@ async def resolve_product_blocks(
         stmt = select_stmt
 
     stmt = sort_product_blocks(stmt, pydantic_sort_by, _error_handler)
-    total = db.session.scalar(select(func.count()).select_from(stmt.subquery()))
+    total = await session.scalar(select(func.count()).select_from(stmt.subquery()))
     stmt = apply_range_to_statement(stmt, after, after + first + 1)
 
     graphql_product_blocks = []

orchestrator/graphql/resolvers/resource_type.py

@@ -1,7 +1,6 @@
 import structlog
 from sqlalchemy import func, select
 
-from orchestrator.db import db
 from orchestrator.db.filters import Filter
 from orchestrator.db.filters.resource_type import (
     RESOURCE_TYPE_TABLE_COLUMN_CLAUSES,
@@ -32,6 +31,7 @@ async def resolve_resource_types(
     query: str | None = None,
 ) -> Connection[ResourceType]:
     _error_handler = create_resolver_error_handler(info)
+    session = info.context.db_session
 
     pydantic_filter_by: list[Filter] = [item.to_pydantic() for item in filter_by] if filter_by else []
     pydantic_sort_by: list[Sort] = [item.to_pydantic() for item in sort_by] if sort_by else []
@@ -54,7 +54,7 @@ async def resolve_resource_types(
         stmt = select_stmt
 
     stmt = sort_resource_types(stmt, pydantic_sort_by, _error_handler)
-    total = db.session.scalar(select(func.count()).select_from(stmt.subquery()))
+    total = await session.scalar(select(func.count()).select_from(stmt.subquery()))
     stmt = apply_range_to_statement(stmt, after, after + first + 1)
 
     graphql_resource_types = []

orchestrator/graphql/resolvers/subscription.py

@@ -21,7 +21,7 @@
 from strawberry.experimental.pydantic.conversion_types import StrawberryTypeFromPydantic
 
 from nwastdlib.asyncio import gather_nice
-from orchestrator.db import ProductTable, SubscriptionTable, db
+from orchestrator.db import ProductTable, SubscriptionTable
 from orchestrator.db.filters import Filter
 from orchestrator.db.filters.subscription import (
     filter_by_query_string,
@@ -99,9 +99,10 @@ async def format_subscription(info: OrchestratorInfo, subscription: Subscription
 
 
 async def resolve_subscription(info: OrchestratorInfo, id: UUID) -> SubscriptionInterface | None:
+    session = info.context.db_session
     stmt = select(SubscriptionTable).where(SubscriptionTable.subscription_id == id)
 
-    if subscription := db.session.scalar(stmt):
+    if subscription := await session.scalar(stmt):
         return await format_subscription(info, subscription)
     return None
 
@@ -115,6 +116,7 @@ async def resolve_subscriptions(
     query: str | None = None,
 ) -> Connection[SubscriptionInterface]:
     _error_handler = create_resolver_error_handler(info)
+    session = info.context.db_session
 
     pydantic_filter_by: list[Filter] = [item.to_pydantic() for item in filter_by] if filter_by else []
     pydantic_sort_by: list[Sort] = [item.to_pydantic() for item in sort_by] if sort_by else []
@@ -141,12 +143,13 @@ async def resolve_subscriptions(
         stmt = filter_by_query_string(stmt, query)
 
     stmt = cast(Select, sort_subscriptions(stmt, pydantic_sort_by, _error_handler))
-    total = db.session.scalar(select(func.count()).select_from(stmt.subquery()))
+    total = await session.scalar(select(func.count()).select_from(stmt.subquery()))
     stmt = apply_range_to_statement(stmt, after, after + first + 1)
 
     graphql_subscriptions: list[SubscriptionInterface] = []
     if is_querying_page_data(info):
-        subscriptions = db.session.scalars(stmt).all()
+        query_result = await session.scalars(stmt)
+        subscriptions = query_result.all()
         graphql_subscriptions = list(await gather_nice((format_subscription(info, p) for p in subscriptions)))  # type: ignore
     logger.info("Resolve subscriptions", filter_by=filter_by, total=total)
 

orchestrator/graphql/resolvers/workflow.py

@@ -1,7 +1,6 @@
 import structlog
 from sqlalchemy import func, select
 
-from orchestrator.db import db
 from orchestrator.db.filters import Filter
 from orchestrator.db.filters.workflow import WORKFLOW_TABLE_COLUMN_CLAUSES, filter_workflows, workflow_filter_fields
 from orchestrator.db.models import WorkflowTable
@@ -28,6 +27,7 @@ async def resolve_workflows(
     query: str | None = None,
 ) -> Connection[Workflow]:
     _error_handler = create_resolver_error_handler(info)
+    session = info.context.db_session
 
     pydantic_filter_by: list[Filter] = [item.to_pydantic() for item in filter_by] if filter_by else []
     pydantic_sort_by: list[Sort] = [item.to_pydantic() for item in sort_by] if sort_by else []
@@ -49,7 +49,7 @@ async def resolve_workflows(
         stmt = select_stmt
 
     stmt = sort_workflows(stmt, pydantic_sort_by, _error_handler)
-    total = db.session.scalar(select(func.count()).select_from(stmt.subquery()))
+    total = await session.scalar(select(func.count()).select_from(stmt.subquery()))
 
     stmt = apply_range_to_statement(stmt, after, after + first + 1)
 

orchestrator/graphql/schema.py

@@ -13,7 +13,7 @@
 from collections.abc import Callable, Iterable
 from http import HTTPStatus
 from pathlib import Path
-from typing import Any, Coroutine, Protocol
+from typing import Any, AsyncIterator, Coroutine, Protocol
 
 import strawberry
 import structlog
@@ -30,6 +30,7 @@
 from nwastdlib.graphql.extensions.error_handler_extension import ErrorHandlerExtension, ErrorType
 from oauth2_lib.fastapi import AuthManager
 from oauth2_lib.strawberry import authenticated_field
+from orchestrator.db import async_db
 from orchestrator.domain.base import SubscriptionModel
 from orchestrator.graphql.autoregistration import create_subscription_strawberry_type, register_domain_models
 from orchestrator.graphql.extensions.model_cache import ModelCacheExtension
@@ -156,11 +157,15 @@ def default_context_getter(
     auth_manager: AuthManager,
     graphql_models: StrawberryModelType,
     broadcast_thread: ProcessDataBroadcastThread | None = None,
-) -> Callable[[], Coroutine[Any, Any, OrchestratorContext]]:
-    async def context_getter() -> OrchestratorContext:
-        return OrchestratorContext(
-            auth_manager=auth_manager, graphql_models=graphql_models, broadcast_thread=broadcast_thread
-        )
+) -> Callable[[], AsyncIterator[OrchestratorContext]]:
+    async def context_getter() -> AsyncIterator[OrchestratorContext]:
+        async with async_db.session() as session:
+            yield OrchestratorContext(
+                auth_manager=auth_manager,
+                graphql_models=graphql_models,
+                broadcast_thread=broadcast_thread,
+                db_session=session,
+            )
 
     return context_getter