Reuse the session between requests

.env.test

@@ -1,3 +1,4 @@
+TEST=True
 REDIS_URL=redis://localhost:6379/0
 DATABASE_URL=postgresql+asyncpg://postgres:postgres@localhost:5432/postgres
 DATABASE_POOL_CLASS=NullPool

app/config.py

@@ -17,6 +17,8 @@ class Settings(BaseSettings):
     REDIS_URL: str = "redis://"
     DATABASE_URL: str = "psql://postgres:"
     DATABASE_POOL_CLASS: str = "AsyncAdaptedQueuePool"
+    DATABASE_POOL_SIZE: int = 10
+    TEST: bool = False
 
 
 settings = Settings()

app/datasources/db/database.py

@@ -1,7 +1,8 @@
 import logging
-from functools import wraps
+from collections.abc import AsyncGenerator
+from functools import cache, wraps
 
-from sqlalchemy.ext.asyncio import create_async_engine
+from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
 from sqlalchemy.pool import AsyncAdaptedQueuePool, NullPool
 from sqlmodel.ext.asyncio.session import AsyncSession
 
@@ -12,15 +13,34 @@
     AsyncAdaptedQueuePool.__name__: AsyncAdaptedQueuePool,
 }
 
-engine = create_async_engine(
-    settings.DATABASE_URL,
-    echo=True,
-    future=True,
-    poolclass=pool_classes.get(settings.DATABASE_POOL_CLASS),
-)
 
-
-def get_database_session(func):
+@cache
+def get_engine() -> AsyncEngine:
+    """
+    Establish connection to database
+    :return:
+    """
+    if settings.TEST:
+        return create_async_engine(
+            settings.DATABASE_URL,
+            future=True,
+            poolclass=NullPool,
+        )
+    else:
+        return create_async_engine(
+            settings.DATABASE_URL,
+            future=True,
+            poolclass=pool_classes.get(settings.DATABASE_POOL_CLASS),
+            pool_size=settings.DATABASE_POOL_SIZE,
+        )
+
+
+async def get_database_session() -> AsyncGenerator:
+    async with AsyncSession(get_engine(), expire_on_commit=False) as session:
+        yield session
+
+
+def database_session(func):
     """
     Decorator that creates a new database session for the given function
 
@@ -30,16 +50,13 @@ def get_database_session(func):
 
     @wraps(func)
     async def wrapper(*args, **kwargs):
-        async with AsyncSession(engine) as session:
+        async with AsyncSession(get_engine(), expire_on_commit=False) as session:
             try:
                 return await func(*args, **kwargs, session=session)
             except Exception as e:
                 # Rollback errors
                 await session.rollback()
                 logging.error(f"Error occurred: {e}")
                 raise
-            finally:
-                # Ensure that session is closed
-                await session.close()
 
     return wrapper

app/routers/contracts.py

@@ -1,7 +1,10 @@
 from typing import Sequence
 
-from fastapi import APIRouter
+from fastapi import APIRouter, Depends
 
+from sqlmodel.ext.asyncio.session import AsyncSession
+
+from ..datasources.db.database import get_database_session
 from ..datasources.db.models import Contract
 from ..services.contract import ContractService
 
@@ -12,6 +15,7 @@
 
 
 @router.get("", response_model=Sequence[Contract])
-async def list_contracts() -> Sequence[Contract]:
-    contract_service = ContractService()
-    return await contract_service.get_all()
+async def list_contracts(
+    session: AsyncSession = Depends(get_database_session),
+) -> Sequence[Contract]:
+    return await ContractService.get_all(session)

app/services/contract.py

@@ -3,14 +3,12 @@
 from sqlmodel import select
 from sqlmodel.ext.asyncio.session import AsyncSession
 
-from app.datasources.db.database import get_database_session
 from app.datasources.db.models import Contract
 
 
 class ContractService:
 
     @staticmethod
-    @get_database_session
     async def get_all(session: AsyncSession) -> Sequence[Contract]:
         """
         Get all contracts
@@ -22,7 +20,6 @@ async def get_all(session: AsyncSession) -> Sequence[Contract]:
         return result.all()
 
     @staticmethod
-    @get_database_session
     async def create(contract: Contract, session: AsyncSession) -> Contract:
         """
         Create a new contract

app/tests/db/db_async_conn.py

@@ -2,12 +2,12 @@
 
 from sqlmodel import SQLModel
 
-from app.datasources.db.database import engine
+from app.datasources.db.database import get_engine
 
 
 class DbAsyncConn(unittest.IsolatedAsyncioTestCase):
     async def asyncSetUp(self):
-        self.engine = engine
+        self.engine = get_engine()
         # Create the database tables
         async with self.engine.begin() as conn:
             await conn.run_sync(SQLModel.metadata.create_all)

app/tests/db/test_model.py

@@ -1,13 +1,13 @@
 from sqlmodel import select
 from sqlmodel.ext.asyncio.session import AsyncSession
 
-from app.datasources.db.database import get_database_session
+from app.datasources.db.database import database_session
 from app.datasources.db.models import Contract
 from app.tests.db.db_async_conn import DbAsyncConn
 
 
 class TestModel(DbAsyncConn):
-    @get_database_session
+    @database_session
     async def test_contract(self, session: AsyncSession):
         contract = Contract(address=b"a", name="A Test Contracts")
         session.add(contract)

app/tests/routers/test_contracts.py

@@ -1,5 +1,8 @@
 from fastapi.testclient import TestClient
 
+from sqlmodel.ext.asyncio.session import AsyncSession
+
+from ...datasources.db.database import database_session
 from ...datasources.db.models import Contract
 from ...main import app
 from ...services.contract import ContractService
@@ -13,14 +16,15 @@ class TestRouterContract(DbAsyncConn):
     def setUpClass(cls):
         cls.client = TestClient(app)
 
-    async def test_view_contracts(self):
+    @database_session
+    async def test_view_contracts(self, session: AsyncSession):
         contract = Contract(address=b"a", name="A Test Contracts")
         expected_response = {
             "name": "A Test Contracts",
             "description": None,
             "address": "a",
         }
-        await ContractService.create(contract=contract)
+        await ContractService.create(contract=contract, session=session)
         response = self.client.get("/api/v1/contracts")
         self.assertEqual(response.status_code, 200)
         self.assertDictEqual(response.json()[0], expected_response)