Feat: Add Report Repository for Async Operations

backend/app/core/config.py

@@ -6,7 +6,7 @@ class Settings(BaseSettings):
     APP_NAME: str = "ChainReport API"
     DEBUG: bool = False
     API_V1_PREFIX: str = "/api/v1"
-    DATABASE_URL: str = "sqlite:///./sql_app.db"
+    DATABASE_URL: str = "sqlite+aiosqlite:///./sql_app.db"
     ONCHAIN_METRICS_URL: str | None = None
     TOKENOMICS_URL: str | None = None
     TWITTER_API_KEY: str = ""
@@ -23,6 +23,20 @@ class Settings(BaseSettings):
     WHITEPAPER_TEXT_SOURCES: Dict[str, str] = {}
     CODE_AUDIT_REPO_URL: str | None = None
 
+    # Database connection settings for PostgreSQL (if used)
+    DB_USER: str | None = None
+    DB_PASSWORD: str | None = None
+    DB_HOST: str | None = None
+    DB_PORT: str | None = None
+    DB_NAME: str | None = None
+
+    # Test database connection settings
+    TEST_DB_USER: str | None = None
+    TEST_DB_PASSWORD: str | None = None
+    TEST_DB_HOST: str | None = None
+    TEST_DB_PORT: str | None = None
+    TEST_DB_NAME: str | None = None
+
     model_config = SettingsConfigDict(env_file=".env")
 
 

backend/app/db/database.py

@@ -1,24 +1,28 @@
-from typing import Generator
-from sqlalchemy import create_engine
-from sqlalchemy.orm import sessionmaker, Session
+from typing import AsyncGenerator
+from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
 from backend.app.db.base import Base
-
+from backend.app.core.config import settings
 
 SQLALCHEMY_DATABASE_URL = settings.DATABASE_URL
 
 if SQLALCHEMY_DATABASE_URL.startswith("sqlite"):
-    engine = create_engine(
+    # SQLite async support via aiosqlite driver
+    engine = create_async_engine(
         SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False}
     )
 else:
-    engine = create_engine(SQLALCHEMY_DATABASE_URL)
-SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+    engine = create_async_engine(SQLALCHEMY_DATABASE_URL)
+
+AsyncSessionLocal = async_sessionmaker(
+    autocommit=False,
+    autoflush=False,
+    bind=engine,
+    class_=AsyncSession,
+    expire_on_commit=False,
+)
+
 
-Base = declarative_base()
 
-def get_db() -> Generator[Session, None, None]:
-    db: Session = SessionLocal()
-    try:
-        yield db
-    finally:
-        db.close()
+async def get_db() -> AsyncGenerator[AsyncSession, None]:
+    async with AsyncSessionLocal() as session:
+        yield session

backend/app/db/repositories/report_repository.py

@@ -0,0 +1,71 @@
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy import select, update
+from sqlalchemy.exc import IntegrityError
+from backend.app.db.models.report import Report
+from backend.app.db.models.report_state import ReportState, ReportStatusEnum
+from typing import Dict, Any
+
+class ReportRepository:
+    def __init__(self, session: AsyncSession):
+        self.session = session
+
+    async def create_report_entry(self, report_id: str) -> Report:
+        try:
+            report = Report(id=report_id)
+            self.session.add(report)
+            report_state = ReportState(report_id=report_id, status=ReportStatusEnum.PENDING)
+            self.session.add(report_state)
+            await self.session.commit()
+            await self.session.refresh(report)
+            return report
+        except IntegrityError:
+            await self.session.rollback()
+            # If a report with this ID already exists, fetch and return it
+            existing_report = await self.session.execute(select(Report).where(Report.id == report_id))
+            report = existing_report.scalar_one_or_none()
+            if report:
+                return report
+            else:
+                # This case should ideally not be reached if IntegrityError is due to report_id
+                raise
+        except Exception:
+            await self.session.rollback()
+            raise
+
+    async def update_report_status(self, report_id: str, status: ReportStatusEnum) -> ReportState | None:
+        try:
+            stmt = update(ReportState).where(ReportState.report_id == report_id).values(status=status).returning(ReportState)
+            result = await self.session.execute(stmt)
+            updated_report_state = result.scalar_one_or_none()
+            await self.session.commit()
+            return updated_report_state
+        except Exception:
+            await self.session.rollback()
+            raise
+
+    async def store_partial_report_results(self, report_id: str, partial_data: Dict[str, Any]) -> ReportState | None:
+        try:
+            stmt = update(ReportState).where(ReportState.report_id == report_id).values(partial_agent_output=partial_data).returning(ReportState)
+            result = await self.session.execute(stmt)
+            updated_report_state = result.scalar_one_or_none()
+            await self.session.commit()
+            return updated_report_state
+        except Exception:
+            await self.session.rollback()
+            raise
+
+    async def save_final_report(self, report_id: str, data: Dict[str, Any]) -> ReportState | None:
+        try:
+            stmt = update(ReportState).where(ReportState.report_id == report_id).values(final_report_json=data, status=ReportStatusEnum.COMPLETED).returning(ReportState)
+            result = await self.session.execute(stmt)
+            updated_report_state = result.scalar_one_or_none()
+            await self.session.commit()
+            return updated_report_state
+        except Exception:
+            await self.session.rollback()
+            raise
+
+    async def get_report_state(self, report_id: str) -> ReportState | None:
+        stmt = select(ReportState).where(ReportState.report_id == report_id)
+        result = await self.session.execute(stmt)
+        return result.scalar_one_or_none()