Add benchmarks to the worker and some fixes to the frontend

Author: pablogsal

backend/app/admin_auth.py

@@ -112,7 +112,23 @@ async def require_admin_auth(
     #     is_active=True
     # )
     # return fake_session
+    if not admin_session_token:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Admin authentication required",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
 
+    try:
+        session = await get_admin_session(db, admin_session_token)
+    except Exception as e:
+        # Log the database error but don't expose internal details
+        logger.error(f"Database error in admin auth: {e}")
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Authentication service unavailable",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
     if not session:
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,

backend/app/models.py

@@ -1,6 +1,7 @@
 from sqlalchemy import (
     Column,
     Integer,
+    BigInteger,
     String,
     DateTime,
     Text,
@@ -118,9 +119,9 @@ class BenchmarkResult(Base):
     id = Column(String(200), primary_key=True)
     run_id = Column(String(100), ForeignKey("runs.run_id"), nullable=False)
     benchmark_name = Column(String(100), nullable=False)
-    high_watermark_bytes = Column(Integer, nullable=False)
+    high_watermark_bytes = Column(BigInteger, nullable=False)
     allocation_histogram = Column(JSON, nullable=False)  # Array of [size, count] tuples
-    total_allocated_bytes = Column(Integer, nullable=False)
+    total_allocated_bytes = Column(BigInteger, nullable=False)
     top_allocating_functions = Column(JSON, nullable=False)  # Array of function objects
     flamegraph_html = Column(Text, nullable=True)  # HTML content of the flamegraph
 

backend/app/routers/admin.py

@@ -1080,6 +1080,9 @@ async def execute_query(
 
     start_time = time.time()
 
+    # Get admin username now while in async context
+    admin_username = admin_session.github_username
+    
     try:
         # Execute the query
         result = await db.execute(text(query))
@@ -1121,7 +1124,7 @@ async def execute_query(
 
         logger.error(f"Query execution failed: {e}", extra={
             "query": query,
-            "admin_user": admin_session.github_username,
+            "admin_user": admin_username,
         })
 
         return QueryResult(

backend/app/routers/benchmarks.py

@@ -73,10 +73,7 @@ async def get_diff_table(
         db, commit_sha=commit_sha, binary_id=binary_id, environment_id=environment_id
     )
     if not runs:
-        raise HTTPException(
-            status_code=404,
-            detail="No runs found for this commit, binary, and environment",
-        )
+        return []
 
     current_run = runs[0]  # Get the latest run
     current_results = await crud.get_benchmark_results(db, run_id=current_run.run_id)

frontend/src/app/admin/components/QueryConsole.tsx

@@ -193,16 +193,36 @@ export default function QueryConsole() {
           });
         }
       } else {
-        const errorData = await response.json();
-        setResult({
-          success: false,
-          error: errorData.detail || 'Query execution failed',
-        });
-        toast({
-          title: 'Error',
-          description: errorData.detail || 'Query execution failed',
-          variant: 'destructive',
-        });
+        // Get response text first, then try to parse as JSON
+        const responseText = await response.text();
+        
+        try {
+          const errorData = JSON.parse(responseText);
+          setResult({
+            success: false,
+            error: errorData.detail || 'Query execution failed',
+          });
+          toast({
+            title: 'Error',
+            description: errorData.detail || 'Query execution failed',
+            variant: 'destructive',
+          });
+        } catch (jsonError) {
+          // Handle non-JSON error responses (e.g., HTML "Internal Server Error")
+          const errorMessage = responseText.includes('Internal Server Error') 
+            ? 'Internal server error occurred. Please check the query syntax and try again.'
+            : responseText || 'Query execution failed';
+          
+          setResult({
+            success: false,
+            error: errorMessage,
+          });
+          toast({
+            title: 'Error',
+            description: errorMessage,
+            variant: 'destructive',
+          });
+        }
       }
     } catch (error) {
       console.error('Error executing query:', error);
@@ -537,8 +557,9 @@ WHERE expires_at < NOW() OR is_active = false;`
 WHERE high_watermark_bytes > 1000000000; -- 1GB threshold`
       },
       {
-        name: '🗑️ Remove old Python version data',
-        query: `DELETE FROM runs 
+        name: '🗑️ Remove old Python version data (Step 1: benchmark_results)',
+        query: `-- First delete benchmark_results for old Python versions
+DELETE FROM benchmark_results 
 WHERE run_id IN (
   SELECT r.run_id 
   FROM runs r 
@@ -547,33 +568,66 @@ WHERE run_id IN (
 );`
       },
       {
-        name: '🗑️ Delete all data for specific binary',
-        query: `-- Delete all runs and results for a specific binary
+        name: '🗑️ Remove old Python version data (Step 2: runs)',
+        query: `-- Then delete runs for old Python versions
+DELETE FROM runs 
+WHERE run_id IN (
+  SELECT r.run_id 
+  FROM runs r 
+  JOIN commits c ON r.commit_sha = c.sha 
+  WHERE c.python_major < 3 OR (c.python_major = 3 AND c.python_minor < 10)
+);`
+      },
+      {
+        name: '🗑️ Delete specific Python version data',
+        query: `-- Template: Replace X with Python minor version (e.g., 15)
+-- Run this query twice: first for benchmark_results, then for runs
 DELETE FROM benchmark_results 
 WHERE run_id IN (
-  SELECT run_id FROM runs WHERE binary_id = 'REPLACE_WITH_BINARY_ID'
+  SELECT run_id FROM runs WHERE python_minor = X
 );
+-- DELETE FROM runs WHERE python_minor = X;`
+      },
+      {
+        name: '🗑️ Delete binary data (Step 1: benchmark_results)',
+        query: `-- First delete benchmark_results for specific binary
+DELETE FROM benchmark_results 
+WHERE run_id IN (
+  SELECT run_id FROM runs WHERE binary_id = 'REPLACE_WITH_BINARY_ID'
+);`
+      },
+      {
+        name: '🗑️ Delete binary data (Step 2: runs)',
+        query: `-- Then delete runs for specific binary
 DELETE FROM runs WHERE binary_id = 'REPLACE_WITH_BINARY_ID';`
       },
       {
-        name: '🗑️ Delete all data for specific environment',
-        query: `-- Delete all runs and results for a specific environment
+        name: '🗑️ Delete environment data (Step 1: benchmark_results)',
+        query: `-- First delete benchmark_results for specific environment
 DELETE FROM benchmark_results 
 WHERE run_id IN (
   SELECT run_id FROM runs WHERE environment_id = 'REPLACE_WITH_ENVIRONMENT_ID'
-);
+);`
+      },
+      {
+        name: '🗑️ Delete environment data (Step 2: runs)',
+        query: `-- Then delete runs for specific environment
 DELETE FROM runs WHERE environment_id = 'REPLACE_WITH_ENVIRONMENT_ID';`
       },
       {
-        name: '🗑️ Delete all data older than X months',
-        query: `-- Delete all benchmark data older than 6 months
+        name: '🗑️ Delete old data (Step 1: benchmark_results)',
+        query: `-- First delete benchmark_results older than 6 months
 DELETE FROM benchmark_results 
 WHERE run_id IN (
   SELECT r.run_id 
   FROM runs r 
   JOIN commits c ON r.commit_sha = c.sha 
   WHERE c.timestamp < NOW() - INTERVAL '6 months'
-);
+);`
+      },
+      {
+        name: '🗑️ Delete old data (Step 2: runs)',
+        query: `-- Then delete runs older than 6 months
 DELETE FROM runs 
 WHERE run_id IN (
   SELECT r.run_id 

frontend/src/components/layout/MemrayStatusIndicator.tsx

@@ -105,8 +105,8 @@ export default function MemrayStatusIndicator() {
                 ))}
               </div>
             ) : (
-              <div className="p-4 bg-green-50 border border-green-200 rounded-lg">
-                <div className="text-green-800 text-center">
+              <div className="p-4 bg-green-50 dark:bg-green-950 border border-green-200 dark:border-green-800 rounded-lg">
+                <div className="text-green-800 dark:text-green-200 text-center">
                   🎉 All memray builds are currently successful across all environments!
                 </div>
               </div>

frontend/src/components/memray/FailureCard.tsx

@@ -54,13 +54,13 @@ export default function FailureCard({
   if (compact) {
     return (
       <>
-        <div className="p-3 border rounded-lg bg-amber-50 border-amber-200">
+        <div className="p-3 border rounded-lg bg-amber-50 dark:bg-amber-950 border-amber-200 dark:border-amber-800">
           <div className="flex items-center justify-between">
             <div className="flex-1">
-              <div className="font-medium text-amber-900">
+              <div className="font-medium text-amber-900 dark:text-amber-200">
                 {binary_name} on {environment_name}
               </div>
-              <div className="text-sm text-amber-700 mt-1">
+              <div className="text-sm text-amber-700 dark:text-amber-300 mt-1">
                 Commit: {commit_sha.substring(0, 8)} • {formatTimestamp(commit_timestamp || failure_timestamp)}
               </div>
             </div>

worker/src/memory_tracker_worker/benchmarks/__init__.py

@@ -110,8 +110,10 @@ def run_benchmarks(venv_dir: Path, output_dir: Path, commit: git.Commit) -> None
     # Create temporary directory for benchmark files
     temp_dir = Path(tempfile.mkdtemp(prefix="benchmarks_"))
     try:
-        # Copy benchmark files to temporary directory
+        # Copy benchmark files and directories to temporary directory
         benchmarks_dir = Path(__file__).parent
+        
+        # Copy all benchmark files
         for benchmark_file in benchmarks_dir.glob("*.py"):
             if benchmark_file.name == "__init__.py":
                 continue
@@ -122,10 +124,24 @@ def run_benchmarks(venv_dir: Path, output_dir: Path, commit: git.Commit) -> None
             logger.info(
                 f"Copied benchmark {benchmark_file.name} to temporary directory"
             )
+        
+        # Copy all data directories
+        for item in benchmarks_dir.iterdir():
+            if item.is_dir():
+                dest_dir = temp_dir / item.name
+                shutil.copytree(item, dest_dir)
+                logger.info(
+                    f"Copied directory {item.name} to temporary directory"
+                )
 
-            # Run benchmark with memray
+        # Run benchmarks with memray
+        for benchmark_file in benchmarks_dir.glob("*.py"):
+            if benchmark_file.name == "__init__.py":
+                continue
+                
             benchmark_name = benchmark_file.stem
             logger.info(f"Running benchmark: {benchmark_name}")
+            dest_file = temp_dir / benchmark_file.name
 
             # Run memray
             memray_output = output_dir / f"{benchmark_name}.bin"