File size: 3,016 Bytes
022fb5a 6ca6336 022fb5a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | import os
from fastapi import APIRouter, HTTPException
from sqlalchemy import text
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PayloadSchemaType
from alembic.config import Config
from alembic import command
from fastapi.concurrency import run_in_threadpool
from ..database import engine
from ..config import get_settings
router = APIRouter()
@router.post("/reset-db")
async def reset_db(password: str = None):
if password != "ketanD048":
raise HTTPException(status_code=403, detail="Invalid reset password.")
settings = get_settings()
try:
# 1. Wipe Postgres natively (drop schema and recreate)
async with engine.begin() as conn:
# We use CASCADE to drop everything including tables, types, etc.
await conn.execute(text('DROP SCHEMA public CASCADE'))
await conn.execute(text('CREATE SCHEMA public'))
await conn.execute(text('GRANT ALL ON SCHEMA public TO public'))
# 2. Wipe Qdrant Collection
q = QdrantClient(url=settings.qdrant_url, api_key=settings.qdrant_api_key)
try:
q.delete_collection(settings.collection_name)
except Exception:
# OK if collection didn't exist
pass
q.create_collection(
collection_name=settings.collection_name,
vectors_config=VectorParams(size=settings.vector_size, distance=Distance.COSINE)
)
# Reinject indices required by the pipeline
q.create_payload_index(
collection_name=settings.collection_name,
field_name="session_id",
field_schema=PayloadSchemaType.KEYWORD
)
q.create_payload_index(
collection_name=settings.collection_name,
field_name="years_of_experience",
field_schema=PayloadSchemaType.FLOAT
)
# 3. Run Alembic Migrations to restore schema
# We assume the API is running from the 'backend' directory where alembic.ini resides.
# If not, we attempt to find it relative to this file.
base_path = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
ini_path = os.path.join(base_path, "alembic.ini")
if not os.path.exists(ini_path):
# Fallback for different envs
ini_path = "alembic.ini"
def _run_migrations():
alembic_cfg = Config(ini_path)
# Ensure the script location is absolute to avoid confusion
alembic_cfg.set_main_option("script_location", os.path.join(base_path, "alembic"))
command.upgrade(alembic_cfg, "head")
await run_in_threadpool(_run_migrations)
return {"status": "success", "message": "Database and Vector Store have been reset and migrations reapplied."}
except Exception as e:
print(f"RESET ERROR: {e}")
raise HTTPException(status_code=500, detail=str(e))
|