""" PolicyEye Database Layer Production: Supabase PostgreSQL (set DATABASE_URL in .env) Dev fallback: SQLite via aiosqlite (auto-created if DATABASE_URL is not set) All queries are fully async via SQLAlchemy ORM. This module exposes the same async helper functions used across the codebase: - init_db() - save_policy(...) - get_policy(policy_id) - get_policy_by_hash(pdf_hash) - get_all_policies() - save_eligibility_check(...) - get_check_history(limit) It uses SQLAlchemy 2.0 async ORM so future migrations with Alembic are straightforward. """ import os import json import logging from typing import List, Optional from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine, AsyncSession from sqlalchemy.orm import declarative_base, sessionmaker, Mapped, mapped_column from sqlalchemy import Integer, String, Float, Text, DateTime, func, select, UniqueConstraint from core.config import DATABASE_PATH logger = logging.getLogger(__name__) # Database URL: prefer DATABASE_URL env var (production), fallback to SQLite (dev) # Production: DATABASE_URL=postgresql+asyncpg://user:pass@host/dbname _raw_db_url = os.getenv("DATABASE_URL") if not _raw_db_url: # Dev fallback: local SQLite file (auto-created, no setup needed) import warnings warnings.warn( "DATABASE_URL not set — using local SQLite. " "Set DATABASE_URL to a PostgreSQL connection string for production.", stacklevel=2, ) sqlite_path = os.path.join(os.path.dirname(__file__), "secureshield_dev.db") DATABASE_URL = f"sqlite+aiosqlite:///{sqlite_path}" else: # Use make_url to properly handle special characters in passwords from sqlalchemy.engine import make_url DATABASE_URL = make_url(_raw_db_url) # Async engine and session # For Supabase Connection Pooler (pgbouncer in transaction mode), # we must disable asyncpg's prepared statement cache. _connect_args = {} if _raw_db_url and "pooler.supabase.com" in _raw_db_url: _connect_args = { "statement_cache_size": 0, "prepared_statement_cache_size": 0, } engine: AsyncEngine = create_async_engine( DATABASE_URL, echo=False, future=True, connect_args=_connect_args ) AsyncSessionLocal = sessionmaker(bind=engine, class_=AsyncSession, expire_on_commit=False) Base = declarative_base() class UserProfile(Base): __tablename__ = "user_profiles" user_id: Mapped[str] = mapped_column(String, primary_key=True) full_name: Mapped[Optional[str]] = mapped_column(String, nullable=True) phone: Mapped[Optional[str]] = mapped_column(String, nullable=True) dob: Mapped[Optional[str]] = mapped_column(String, nullable=True) address: Mapped[Optional[str]] = mapped_column(Text, nullable=True) avatar_base64: Mapped[Optional[str]] = mapped_column(Text, nullable=True) updated_at: Mapped[str] = mapped_column(DateTime(timezone=True), server_default=func.now(), onupdate=func.now()) class Policy(Base): __tablename__ = "policies" __table_args__ = ( UniqueConstraint("user_id", "raw_text_hash", name="uq_policy_user_hash"), ) id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) user_id: Mapped[str] = mapped_column(String, nullable=False, index=True) insurer: Mapped[str] = mapped_column(String, nullable=False) plan_name: Mapped[str] = mapped_column(String, nullable=False) sum_insured: Mapped[float] = mapped_column(Float, nullable=False) policy_type: Mapped[str] = mapped_column(String, default="individual") rules_json: Mapped[str] = mapped_column(Text, nullable=False) raw_text_hash: Mapped[Optional[str]] = mapped_column(String, index=True, nullable=True) pdf_storage_url: Mapped[Optional[str]] = mapped_column(String, nullable=True) is_reviewed: Mapped[int] = mapped_column(Integer, default=0) created_at: Mapped[str] = mapped_column(DateTime(timezone=True), server_default=func.now()) class EligibilityCheck(Base): __tablename__ = "eligibility_checks" id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) user_id: Mapped[str] = mapped_column(String, nullable=False, index=True) policy_id: Mapped[int] = mapped_column(Integer) case_json: Mapped[str] = mapped_column(Text, nullable=False) verdict_json: Mapped[str] = mapped_column(Text, nullable=False) explanation: Mapped[Optional[str]] = mapped_column(Text, nullable=True) created_at: Mapped[str] = mapped_column(DateTime(timezone=True), server_default=func.now()) class ChatThread(Base): __tablename__ = "chat_threads" id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) user_id: Mapped[str] = mapped_column(String, nullable=False, index=True) title: Mapped[str] = mapped_column(String, nullable=False) created_at: Mapped[str] = mapped_column(DateTime(timezone=True), server_default=func.now()) class ChatMessage(Base): __tablename__ = "chat_messages" id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) thread_id: Mapped[int] = mapped_column(Integer, nullable=False, index=True) role: Mapped[str] = mapped_column(String, nullable=False) # 'user' or 'assistant' content: Mapped[str] = mapped_column(Text, nullable=False) method: Mapped[Optional[str]] = mapped_column(String, nullable=True) duration_ms: Mapped[Optional[float]] = mapped_column(Float, nullable=True) created_at: Mapped[str] = mapped_column(DateTime(timezone=True), server_default=func.now()) class AppFeedback(Base): __tablename__ = "app_feedback" id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) user_id: Mapped[Optional[str]] = mapped_column(String, index=True, nullable=True) accuracy_rating: Mapped[int] = mapped_column(Integer, nullable=False) claim_rejections: Mapped[Optional[str]] = mapped_column(Text, nullable=True) problem_solving: Mapped[Optional[str]] = mapped_column(Text, nullable=True) addon_suggestions: Mapped[Optional[str]] = mapped_column(Text, nullable=True) knowledge_source: Mapped[Optional[str]] = mapped_column(Text, nullable=True) govt_regulations: Mapped[Optional[str]] = mapped_column(Text, nullable=True) developer_centric: Mapped[Optional[str]] = mapped_column(Text, nullable=True) other_suggestions: Mapped[Optional[str]] = mapped_column(Text, nullable=True) created_at: Mapped[str] = mapped_column(DateTime(timezone=True), server_default=func.now()) from sqlalchemy.exc import IntegrityError, ProgrammingError async def init_db(): """Create tables if they don't exist (async).""" try: async with engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) logger.info("[Database] Async ORM tables initialized") except (IntegrityError, ProgrammingError) as e: logger.warning(f"[Database] Race condition during table creation, ignoring: {e}") async def save_policy(insurer: str, plan_name: str, sum_insured: float, policy_type: str, rules: List[dict], raw_text_hash: Optional[str], pdf_storage_url: Optional[str] = None, user_id: str = "") -> int: """Save an ingested policy and return its ID.""" async with AsyncSessionLocal() as session: policy = Policy( user_id=user_id, insurer=insurer, plan_name=plan_name, sum_insured=sum_insured, policy_type=policy_type, rules_json=json.dumps(rules), raw_text_hash=raw_text_hash, pdf_storage_url=pdf_storage_url, ) session.add(policy) await session.commit() await session.refresh(policy) logger.info(f"[Database] Saved policy #{policy.id}: {insurer} - {plan_name} (user: {user_id[:8]}...)") return policy.id async def save_user_profile(user_id: str, full_name: str, phone: str, dob: str, address: str, avatar_base64: str) -> bool: """Upsert user profile.""" async with AsyncSessionLocal() as session: stmt = select(UserProfile).where(UserProfile.user_id == user_id).limit(1) result = await session.execute(stmt) profile = result.scalar_one_or_none() if profile: profile.full_name = full_name profile.phone = phone profile.dob = dob profile.address = address if avatar_base64 is not None: profile.avatar_base64 = avatar_base64 else: profile = UserProfile( user_id=user_id, full_name=full_name, phone=phone, dob=dob, address=address, avatar_base64=avatar_base64 ) session.add(profile) await session.commit() logger.info(f"[Database] Saved user profile for {user_id[:8]}...") return True async def get_user_profile(user_id: str) -> Optional[dict]: """Retrieve user profile by ID.""" async with AsyncSessionLocal() as session: stmt = select(UserProfile).where(UserProfile.user_id == user_id).limit(1) result = await session.execute(stmt) profile = result.scalar_one_or_none() if not profile: return None return { "user_id": profile.user_id, "full_name": profile.full_name, "phone": profile.phone, "dob": profile.dob, "address": profile.address, "avatar_base64": profile.avatar_base64, } async def get_policy(policy_id: int, user_id: str = "") -> Optional[dict]: """Retrieve a policy by ID, scoped to user if user_id is provided.""" async with AsyncSessionLocal() as session: if user_id: stmt = select(Policy).where(Policy.id == policy_id, Policy.user_id == user_id).limit(1) result = await session.execute(stmt) policy = result.scalar_one_or_none() else: policy = await session.get(Policy, policy_id) if not policy: return None return { "id": policy.id, "user_id": policy.user_id, "insurer": policy.insurer, "plan_name": policy.plan_name, "sum_insured": policy.sum_insured, "policy_type": policy.policy_type, "rules": json.loads(policy.rules_json), "raw_text_hash": policy.raw_text_hash, "pdf_storage_url": policy.pdf_storage_url, "is_reviewed": policy.is_reviewed, "created_at": str(policy.created_at), } async def get_policy_by_hash(pdf_hash: str, user_id: str = "") -> Optional[dict]: """Retrieve a policy by PDF hash, optionally scoped to a specific user.""" async with AsyncSessionLocal() as session: stmt = select(Policy).where(Policy.raw_text_hash == pdf_hash) if user_id: stmt = stmt.where(Policy.user_id == user_id) result = await session.execute(stmt.limit(1)) policy_obj = result.scalar_one_or_none() if not policy_obj: return None return { "id": policy_obj.id, "user_id": policy_obj.user_id, "insurer": policy_obj.insurer, "plan_name": policy_obj.plan_name, "sum_insured": policy_obj.sum_insured, "policy_type": policy_obj.policy_type, "rules": json.loads(policy_obj.rules_json or "[]"), "raw_text_hash": policy_obj.raw_text_hash, "pdf_storage_url": policy_obj.pdf_storage_url, "is_reviewed": policy_obj.is_reviewed, "created_at": str(policy_obj.created_at), } async def get_all_policies(user_id: str = "") -> List[dict]: """Retrieve all ingested policies for a specific user (without full rules for listing).""" async with AsyncSessionLocal() as session: stmt = select( Policy.id, Policy.insurer, Policy.plan_name, Policy.sum_insured, Policy.policy_type, Policy.is_reviewed, Policy.created_at, ).order_by(Policy.created_at.desc()) if user_id: from sqlalchemy import or_ stmt = stmt.where(or_(Policy.user_id == user_id, Policy.user_id == "api_client")) q = await session.execute(stmt) rows = q.fetchall() return [ { "id": r.id, "insurer": r.insurer, "plan_name": r.plan_name, "sum_insured": r.sum_insured, "policy_type": r.policy_type, "is_reviewed": r.is_reviewed, "created_at": str(r.created_at), } for r in rows ] async def save_eligibility_check(policy_id: int, case_json: str, verdict_json: str, explanation: Optional[str], user_id: str = "") -> int: """Save an eligibility check result.""" async with AsyncSessionLocal() as session: check = EligibilityCheck( user_id=user_id, policy_id=policy_id, case_json=case_json, verdict_json=verdict_json, explanation=explanation, ) session.add(check) await session.commit() await session.refresh(check) return check.id async def get_check_history(limit: int = 20, user_id: str = "") -> List[dict]: """Retrieve recent eligibility check history for a specific user.""" async with AsyncSessionLocal() as session: stmt = select( EligibilityCheck.id, EligibilityCheck.policy_id, EligibilityCheck.case_json, EligibilityCheck.verdict_json, EligibilityCheck.explanation, EligibilityCheck.created_at, ).order_by(EligibilityCheck.created_at.desc()).limit(limit) if user_id: stmt = stmt.where(EligibilityCheck.user_id == user_id) q = await session.execute(stmt) rows = q.fetchall() return [ { "id": r.id, "policy_id": r.policy_id, "case_json": r.case_json, "verdict_json": r.verdict_json, "explanation": r.explanation, "created_at": str(r.created_at), } for r in rows ] async def create_chat_thread(user_id: str, title: str) -> int: """Create a new chat thread.""" async with AsyncSessionLocal() as session: thread = ChatThread(user_id=user_id, title=title) session.add(thread) await session.commit() await session.refresh(thread) return thread.id async def get_chat_threads(user_id: str) -> List[dict]: """Get all chat threads for a user.""" async with engine.begin() as conn: result = await conn.execute( select(ChatThread) .where(ChatThread.user_id == user_id) .order_by(ChatThread.created_at.desc()) ) rows = result.fetchall() return [{"id": r.id, "title": r.title, "created_at": str(r.created_at)} for r in rows] async def delete_chat_thread(thread_id: int): """Delete a chat thread and its messages.""" from sqlalchemy import delete async with engine.begin() as conn: await conn.execute( delete(ChatThread) .where(ChatThread.id == thread_id) ) async def save_chat_message(thread_id: int, role: str, content: str, method: Optional[str] = None, duration_ms: Optional[float] = None) -> int: """Save a chat message to a thread.""" async with AsyncSessionLocal() as session: msg = ChatMessage(thread_id=thread_id, role=role, content=content, method=method, duration_ms=duration_ms) session.add(msg) await session.commit() await session.refresh(msg) return msg.id async def get_chat_messages(thread_id: int) -> List[dict]: """Get all messages for a thread.""" async with AsyncSessionLocal() as session: stmt = select(ChatMessage).where(ChatMessage.thread_id == thread_id).order_by(ChatMessage.created_at.asc()) q = await session.execute(stmt) rows = q.scalars().all() return [{ "id": r.id, "role": r.role, "content": r.content, "method": r.method, "duration_ms": r.duration_ms, "created_at": str(r.created_at) } for r in rows] async def clear_policies_and_checks() -> None: """Delete all rows from `eligibility_checks` and `policies` tables. Useful for tests and demo reset. Uses SQLAlchemy ORM (database-agnostic). """ from sqlalchemy import delete as sa_delete async with AsyncSessionLocal() as session: await session.execute(sa_delete(EligibilityCheck)) await session.execute(sa_delete(Policy)) await session.commit() logger.info("[Database] Cleared all policies and checks.") async def delete_user_data(user_id: str) -> dict: """ Delete all data associated with a user_id to comply with DPDPA 2023 'Right to Erasure' for sensitive health data. """ from sqlalchemy import delete as sa_delete async with AsyncSessionLocal() as session: # Delete profile profile_res = await session.execute(sa_delete(UserProfile).where(UserProfile.user_id == user_id)) # Delete eligibility checks checks_res = await session.execute(sa_delete(EligibilityCheck).where(EligibilityCheck.user_id == user_id)) # Delete policies # Note: We keep the LLM cache (llm_cache table) since it's anonymous and hashed policies_res = await session.execute(sa_delete(Policy).where(Policy.user_id == user_id)) # Delete chat threads (cascade should delete messages, but let's be explicit if needed) threads_res = await session.execute(sa_delete(ChatThread).where(ChatThread.user_id == user_id)) await session.commit() counts = { "profile": profile_res.rowcount, "checks": checks_res.rowcount, "policies": policies_res.rowcount, "chat_threads": threads_res.rowcount, } logger.info(f"[Database] Deleted user data for {user_id}: {counts}") return counts async def save_app_feedback( user_id: str, accuracy_rating: int, claim_rejections: str, problem_solving: str, addon_suggestions: str, knowledge_source: str, govt_regulations: str, developer_centric: str, other_suggestions: str = "" ) -> int: """Save user app feedback.""" async with AsyncSessionLocal() as session: feedback = AppFeedback( user_id=user_id or None, accuracy_rating=accuracy_rating, claim_rejections=claim_rejections, problem_solving=problem_solving, addon_suggestions=addon_suggestions, knowledge_source=knowledge_source, govt_regulations=govt_regulations, developer_centric=developer_centric, other_suggestions=other_suggestions ) session.add(feedback) await session.commit() await session.refresh(feedback) logger.info(f"[Database] Saved app feedback from user: {user_id}") return feedback.id