Add duplicate run prevention and improve database schema

- Replace Run table index with unique constraint on (commit_sha, binary_id, environment_id)
- Add IntegrityError handling in upload endpoint to prevent duplicate runs
- Remove redundant database indexes to improve performance
- Enhance init_db.py with better error handling and database URL override support
- Improve PostgreSQL compatibility with proper table drop ordering
- Add flamegraph cleanup functionality to manage storage efficiently

Author: pablogsal

backend/app/models.py

@@ -8,6 +8,7 @@
     JSON,
     Index,
     Boolean,
+    UniqueConstraint,
 )
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import relationship
@@ -88,8 +89,9 @@ class Run(Base):
     benchmark_results = relationship("BenchmarkResult", back_populates="run")
 
     __table_args__ = (
-        Index(
-            "idx_runs_commit_binary_env", "commit_sha", "binary_id", "environment_id"
+        UniqueConstraint(
+            "commit_sha", "binary_id", "environment_id", 
+            name="unique_commit_binary_env"
         ),
         Index("idx_runs_timestamp", "timestamp"),
         Index(
@@ -101,20 +103,13 @@ class Run(Base):
         Index(
             "idx_runs_env_timestamp", "environment_id", "timestamp"
         ),  # For environment-based queries
-        Index(
-            "idx_runs_timestamp_desc",
-            "timestamp",
-            postgresql_using="btree",
-            mysql_length={"timestamp": 255},
-        ),  # For ORDER BY timestamp DESC
         Index(
             "idx_runs_env_python_timestamp",
             "environment_id",
             "python_major",
             "python_minor",
             "timestamp",
         ),  # For filtered queries
-        Index("idx_runs_binary_env_commit", "binary_id", "environment_id", "commit_sha"),
     )
 
 

backend/app/routers/upload.py

@@ -2,6 +2,8 @@
 
 from fastapi import APIRouter, Depends, HTTPException
 from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy import select, delete, desc, func, update, and_
+from sqlalchemy.exc import IntegrityError
 from datetime import datetime
 import logging
 
@@ -12,6 +14,81 @@
 router = APIRouter(prefix="/api", tags=["upload"])
 
 
+async def cleanup_old_flamegraphs_if_needed(db: AsyncSession, environment_id: str, max_flamegraphs: int = 100):
+    """
+    Set flamegraph_html to NULL for runs older than the last max_flamegraphs in this environment.
+    """
+    logger = logging.getLogger(__name__)
+    
+    keep_runs_result = await db.execute(
+        select(models.Run.run_id).where(
+            models.Run.environment_id == environment_id
+        ).order_by(desc(models.Run.timestamp)).limit(max_flamegraphs)
+    )
+    keep_run_ids = [row[0] for row in keep_runs_result.fetchall()]
+    
+    if not keep_run_ids:
+        return 0
+    
+    total_flamegraphs_result = await db.execute(
+        select(func.count()).select_from(models.BenchmarkResult).join(models.Run).where(
+            models.Run.environment_id == environment_id,
+            models.BenchmarkResult.flamegraph_html.is_not(None)
+        )
+    )
+    total_flamegraphs_before = total_flamegraphs_result.scalar()
+    
+    count_query = select(func.count()).select_from(models.BenchmarkResult).where(
+        and_(
+            models.BenchmarkResult.run_id.not_in(keep_run_ids),
+            models.BenchmarkResult.flamegraph_html.is_not(None)
+        )
+    )
+    count_result = await db.execute(count_query)
+    rows_to_clean = count_result.scalar()
+    
+    if rows_to_clean == 0:
+        logger.info(f"No flamegraphs need cleaning for environment '{environment_id}' (all {total_flamegraphs_before} flamegraphs are from recent runs)")
+        return 0
+    
+    cleanup_query = update(models.BenchmarkResult).where(
+        and_(
+            models.BenchmarkResult.run_id.not_in(keep_run_ids),
+            models.BenchmarkResult.flamegraph_html.is_not(None)
+        )
+    ).values(flamegraph_html=None)
+    
+    try:
+        result = await db.execute(cleanup_query)
+        await db.commit()
+        
+        verify_result = await db.execute(count_query)
+        remaining = verify_result.scalar()
+        actual_cleaned = rows_to_clean - remaining
+        
+        final_flamegraphs_result = await db.execute(
+            select(func.count()).select_from(models.BenchmarkResult).join(models.Run).where(
+                models.Run.environment_id == environment_id,
+                models.BenchmarkResult.flamegraph_html.is_not(None)
+            )
+        )
+        total_flamegraphs_after = final_flamegraphs_result.scalar()
+        
+        if actual_cleaned > 0:
+            logger.info(f"Cleaned up {actual_cleaned} flamegraphs for environment '{environment_id}': {total_flamegraphs_before} → {total_flamegraphs_after} flamegraphs remaining")
+        else:
+            logger.error(f"Cleanup FAILED for environment '{environment_id}'. Expected to clean {rows_to_clean}, but {remaining} still remain")
+        
+        return actual_cleaned
+        
+    except Exception as e:
+        logger.error(f"Exception during cleanup for environment '{environment_id}': {e}")
+        await db.rollback()
+        raise
+
+
+
+
 @router.post("/upload-run", response_model=dict)
 async def upload_worker_run(
     upload_data: schemas.WorkerRunUpload,
@@ -106,6 +183,8 @@ async def upload_worker_run(
         )
     logger.info(f"Configure flags validation passed for binary '{binary_id}'")
 
+    # Note: Duplicate commits are now prevented by database unique constraint on (commit_sha, binary_id, environment_id)
+
     # Create or get commit
     logger.debug(f"Looking up commit {commit_sha[:8]} in database")
     commit = await crud.get_commit_by_sha(db, sha=commit_sha)
@@ -203,6 +282,9 @@ async def upload_worker_run(
             f"Upload completed successfully: run_id={run_id}, created {len(created_results)} benchmark results"
         )
 
+        # Clean up old flamegraphs if we have more than 100 runs for this environment
+        await cleanup_old_flamegraphs_if_needed(db, environment_id, max_flamegraphs=100)
+
         return {
             "message": "Worker run uploaded successfully",
             "run_id": run_id,
@@ -213,6 +295,17 @@ async def upload_worker_run(
             "result_ids": created_results,
         }
 
+    except IntegrityError as e:
+        # Handle unique constraint violation for duplicate commit+binary+environment
+        if "unique_commit_binary_env" in str(e).lower():
+            logger.error(f"Upload failed: Duplicate run for commit {commit_sha[:8]}, binary '{binary_id}', environment '{environment_id}'")
+            raise HTTPException(
+                status_code=409,
+                detail=f"A run already exists for commit {commit_sha[:8]} with binary '{binary_id}' and environment '{environment_id}'. Duplicate uploads are not allowed."
+            )
+        else:
+            logger.error(f"Database integrity error during upload: {e}")
+            raise HTTPException(status_code=500, detail=f"Database integrity error: {str(e)}")
     except HTTPException:
         # Re-raise HTTP exceptions (validation errors) as-is
         raise

backend/scripts/init_db.py

@@ -55,21 +55,36 @@ async def init_admin_data():
 async def init_database():
     """Initialize the database by creating all tables."""
     print("Initializing database...")
+    
+    settings = get_settings()
 
     try:
         # Import models to ensure they're registered
-        from app.models import AdminUser, AdminSession, AuthToken, Commit, Binary, Environment, Run, BenchmarkResult
+        from app.models import AdminUser, AdminSession, AuthToken, Commit, Binary, Environment, Run, BenchmarkResult, Base
+        from app.database import create_database_engine
+        
+        # Create a fresh engine with current settings
+        engine = create_database_engine()
+        
+        print(f"Connected to: {settings.database_url.split('@')[-1] if '@' in settings.database_url else settings.database_url}")
 
-        await create_tables()
+        # Create all tables
+        async with engine.begin() as conn:
+            await conn.run_sync(Base.metadata.create_all)
         print("✅ Database tables created successfully!")
 
+        # Close the engine we created
+        await engine.dispose()
+        
         # Initialize admin data
         admin_success = await init_admin_data()
         if not admin_success:
             return False
 
     except Exception as e:
         print(f"❌ Error creating database tables: {e}")
+        import traceback
+        traceback.print_exc()
         return False
 
     return True
@@ -78,24 +93,62 @@ async def init_database():
 async def reset_database():
     """Reset the database by dropping and recreating all tables."""
     print("Resetting database...")
+    
+    settings = get_settings()
+    is_postgres = "postgresql" in settings.database_url
 
     try:
         # Import models to ensure they're registered
-        from app.models import AdminUser, AdminSession, AuthToken, Commit, Binary, Environment, Run, BenchmarkResult
+        from app.models import AdminUser, AdminSession, AuthToken, Commit, Binary, Environment, Run, BenchmarkResult, Base
+        from app.database import create_database_engine
+        from sqlalchemy.ext.asyncio import create_async_engine
 
-        await drop_tables()
-        print("🗑️  Existing tables dropped")
-
-        await create_tables()
+        # Create a fresh engine with current settings
+        engine = create_database_engine()
+        
+        print(f"Connected to: {settings.database_url.split('@')[-1] if '@' in settings.database_url else settings.database_url}")
+        
+        # For PostgreSQL, we need to handle foreign key constraints more carefully
+        if is_postgres:
+            async with engine.begin() as conn:
+                # First, drop all dependent tables (those with foreign keys)
+                print("🗑️  Dropping dependent tables...")
+                await conn.run_sync(BenchmarkResult.__table__.drop, checkfirst=True)
+                await conn.run_sync(Run.__table__.drop, checkfirst=True)
+                await conn.run_sync(AdminSession.__table__.drop, checkfirst=True)
+                
+                # Then drop the referenced tables
+                print("🗑️  Dropping referenced tables...")
+                await conn.run_sync(Commit.__table__.drop, checkfirst=True)
+                await conn.run_sync(Binary.__table__.drop, checkfirst=True)
+                await conn.run_sync(Environment.__table__.drop, checkfirst=True)
+                await conn.run_sync(AuthToken.__table__.drop, checkfirst=True)
+                await conn.run_sync(AdminUser.__table__.drop, checkfirst=True)
+                
+                print("🗑️  All tables dropped")
+        else:
+            # For SQLite, drop all tables using metadata
+            async with engine.begin() as conn:
+                await conn.run_sync(Base.metadata.drop_all)
+            print("🗑️  Existing tables dropped")
+
+        # Recreate all tables
+        async with engine.begin() as conn:
+            await conn.run_sync(Base.metadata.create_all)
         print("✅ Database tables recreated successfully!")
 
+        # Close the engine we created
+        await engine.dispose()
+        
         # Initialize admin data after reset
         admin_success = await init_admin_data()
         if not admin_success:
             return False
 
     except Exception as e:
         print(f"❌ Error resetting database: {e}")
+        import traceback
+        traceback.print_exc()
         return False
 
     return True
@@ -110,9 +163,28 @@ async def reset_database():
         action="store_true",
         help="Reset the database (drop and recreate tables)",
     )
+    parser.add_argument(
+        "--database-url",
+        type=str,
+        help="Database URL to use (overrides DATABASE_URL env var)",
+    )
 
     args = parser.parse_args()
 
+    # Override database URL if provided
+    if args.database_url:
+        os.environ["DATABASE_URL"] = args.database_url
+        # Also try the lowercase version
+        os.environ["database_url"] = args.database_url
+        # Clear the settings cache so it picks up the new DATABASE_URL
+        from app.config import get_settings
+        get_settings.cache_clear()
+        
+        # Verify the database URL is actually being used
+        settings = get_settings()
+        print(f"Using database URL: {args.database_url.split('@')[-1] if '@' in args.database_url else args.database_url}")
+        print(f"Settings database URL: {settings.database_url.split('@')[-1] if '@' in settings.database_url else settings.database_url}")
+
     if args.reset:
         success = asyncio.run(reset_database())
     else: