better handling of migrations and updating of endpoints

Author: Khauneesh-AI

app/main.py

@@ -239,11 +239,25 @@ class StudioUpgradeResponse(BaseModel):
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     """Lifespan context manager for the FastAPI application"""
-    # Create document upload directory on startup
-    #path_manager.make_dirs(path_manager.upload_dir)
+    # Startup code
     os.makedirs(UPLOAD_DIR, exist_ok=True)
     print(f"Document upload directory created at: {UPLOAD_DIR}")
+    
+    # Run database migrations (init_db() already called by DatabaseManager)
+    try:
+        print("Running database migrations via external script...")
+        result = subprocess.run(
+            ["uv", "run", "python", "run_migrations.py"],
+            capture_output=True,
+            text=True,
+            check=True
+        )
+        print(f"Migration completed: {result.stdout}")
+    except subprocess.CalledProcessError as e:
+        print(f"Migration warning: {e.stderr}")
+    
     yield
+    print("Application shutting down...")
 
 
 app = FastAPI(
@@ -338,12 +352,12 @@ def get_timeout_for_request(request: Request) -> float:
 
 
 
-@app.on_event("startup")
-async def startup_event():
-    """Check for and apply any pending migrations on startup"""
-    success, message = await alembic_manager.handle_database_upgrade()
-    if not success:
-        print(f"Warning: {message}")
+# @app.on_event("startup")
+# async def startup_event():
+#     """Check for and apply any pending migrations on startup"""
+#     success, message = await alembic_manager.handle_database_upgrade()
+#     if not success:
+#         print(f"Warning: {message}")
 
 @app.post("/get_project_files", include_in_schema=True, responses = responses, 
            description = "get project file details")
@@ -1318,12 +1332,10 @@ async def get_example_payloads(use_case:UseCase):
     return payload
 
 
-# Add these two endpoints
 @app.get("/synthesis-studio/check-upgrade", response_model=StudioUpgradeStatus)
 async def check_upgrade_status():
     """Check if any upgrades are available"""
     try:
-    
         # Fetch latest changes
         subprocess.run(["git", "fetch"], check=True, capture_output=True)
 
@@ -1357,39 +1369,22 @@ async def check_upgrade_status():
             git_remote_commit=remote_commit,
             updates_available=updates_available
         )
+        
+    except subprocess.CalledProcessError as e:
+        raise HTTPException(status_code=500, detail=f"Git error: {str(e)}")
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
-    
 
-# Ensure uv is installed
-def ensure_uv_installed():
-    try:
-        print(subprocess.run(["uv", "--version"], check=True, capture_output=True))
-    except (subprocess.CalledProcessError, FileNotFoundError):
-        print(subprocess.run(["curl -LsSf https://astral.sh/uv/install.sh | sh"], shell=True, check=True))
-        os.environ["PATH"] = f"{os.environ['HOME']}/.cargo/bin:{os.environ['PATH']}"
 
-# Setup virtual environment and dependencies
-def setup_environment(PROJECT_ROOT):
-    
-    venv_dir = PROJECT_ROOT / ".venv"
-    
-    # Create venv if it doesn't exist - specifying the path explicitly
-    if not venv_dir.exists():
-        print(subprocess.run(["uv", "venv", ".venv"], cwd=PROJECT_ROOT, check=True))
-    
-    # Install dependencies with uv pip instead of sync to avoid pyproject.toml parsing issues
-    print(subprocess.run(["uv", "pip", "install", "-e", "."], cwd=PROJECT_ROOT, check=True))
 
 @app.post("/synthesis-studio/upgrade", response_model=StudioUpgradeResponse)
 async def perform_upgrade():
     """
     Perform upgrade process:
     1. Pull latest code
-    2. Run database migrations with Alembic
-    3. Run build_client.sh
-    4. Run start_application.py
-    5. Restart CML application
+    2. Run database migrations
+    3. Build frontend (includes dependency sync)
+    4. Restart application
     """
     try:
         messages = []
@@ -1398,75 +1393,65 @@ async def perform_upgrade():
         db_upgraded = False
 
         PROJECT_ROOT = Path(os.getcwd())
+        
         # 1. Git operations
         try:
-            
             # Stash any changes
-            print(subprocess.run(["git", "stash"], check=True, capture_output=True))
+            subprocess.run(["git", "stash"], check=True, capture_output=True)
 
             # Pull updates
-            print(subprocess.run(["git", "pull"], check=True, capture_output=True))
+            subprocess.run(["git", "pull"], check=True, capture_output=True)
 
             # Try to pop stash
             try:
-                print(subprocess.run(["git", "stash", "pop"], check=True, capture_output=True))
+                subprocess.run(["git", "stash", "pop"], check=True, capture_output=True)
             except subprocess.CalledProcessError:
                 messages.append("Warning: Could not restore local changes")
 
             git_updated = True
             messages.append("Git repository updated")
-            print(messages)
 
         except subprocess.CalledProcessError as e:
             messages.append(f"Git update failed: {e}")
             raise HTTPException(status_code=500, detail=str(e))
 
         # 2. Database migrations
         try:
-            # In your upgrade endpoint, you can add this debug line:
-            print(f"Current working directory: {os.getcwd()}")
-            print(f"Alembic.ini exists: {os.path.exists('alembic.ini')}")
-            print("--- Starting database migration via external script ---")
-            # Use `uv run` to ensure the script runs within the project's virtual environment
-            # This is more robust than just calling 'python'
             result = subprocess.run(
                 ["uv", "run", "python", "run_migrations.py"],
                 capture_output=True,
                 text=True,
-                check=True  # This will raise CalledProcessError on failure
+                check=True,
+                cwd=PROJECT_ROOT
             )
 
-            print(result.stdout) # Log the output from the script
+            print(result.stdout)
             db_upgraded = True
-            messages.append("Database migration check completed successfully.")
-        except Exception as e:
+            messages.append("Database migrations completed")
+        except subprocess.CalledProcessError as e:
             messages.append(f"Database migration failed: {str(e)}")
             raise HTTPException(status_code=500, detail=str(e))
 
-        # 3. Run build_client.sh
+        # 3. Build frontend (build_client.sh handles uv and dependencies)
         try:
-            
-            ensure_uv_installed()
-            
-            setup_environment(PROJECT_ROOT)
-
-            subprocess.run(["bash build/shell_scripts/build_client.sh"], shell=True, check=True)
+            subprocess.run(
+                ["bash", "build/shell_scripts/build_client.sh"], 
+                check=True,
+                cwd=PROJECT_ROOT
+            )
             frontend_rebuilt = True
             messages.append("Frontend rebuilt successfully")
         except subprocess.CalledProcessError as e:
             messages.append(f"Frontend build failed: {e}")
             raise HTTPException(status_code=500, detail=str(e))
 
-       
+        # 4. Restart application
         if git_updated or frontend_rebuilt or db_upgraded:
             try:
-                # Small delay to ensure logs are captured
-                time.sleep(10)
-                print("application restart will happen now")
+                time.sleep(5)
                 restart_application()
                 messages.append("Application restart initiated")
 
-                # Note: This response might not reach the client due to the restart
                 return StudioUpgradeResponse(
                     success=True,
                     message="; ".join(messages),
@@ -1479,10 +1464,7 @@ async def perform_upgrade():
                 raise HTTPException(status_code=500, detail=str(e))
 
     except Exception as e:
-        raise HTTPException(
-            status_code=500,
-            detail=f"Upgrade failed: {str(e)}"
-        )
+        raise HTTPException(status_code=500, detail=f"Upgrade failed: {str(e)}")
 #****** comment below for testing just backend**************
 current_directory = os.path.dirname(os.path.abspath(__file__))
 client_build_path = os.path.join(current_directory, "client", "dist")

run_migrations.py

@@ -1,33 +1,61 @@
 import asyncio
 import sys
 from pathlib import Path
-
+import subprocess
+import os
 # Ensure the 'app' directory is in the Python path
 ROOT_DIR = Path(__file__).parent
 APP_DIR = ROOT_DIR / "app"
 sys.path.append(str(ROOT_DIR))
 
 from app.migrations.alembic_manager import AlembicMigrationManager
 
+
+
+# Add this function before main()
+def handle_fresh_installation(db_path):
+    """Handle installations where init_db() already created schema"""
+    if os.path.exists(db_path):
+        # Database exists (created by init_db), check migration status
+        try:
+            result = subprocess.run(["alembic", "current"], capture_output=True, text=True, check=True)
+            print(f"Checking migration status: {result.stdout.strip()}")
+            
+            # Check if no migration history exists
+            if not any(len(line.strip()) >= 12 and not line.startswith('INFO') 
+                      for line in result.stdout.split('\n')):
+                print("Database exists (via init_db) but no migration history - stamping with head")
+                subprocess.run(["alembic", "stamp", "head"], check=True)
+                return True
+        except subprocess.CalledProcessError as e:
+            print(f"Error checking migration status: {e.stderr}")
+    return False
+
+
+# Update main() function
 async def main():
     """
-    Initializes the migration manager and runs the database upgrade.
-    This will always use the latest code from disk.
+    Migration script that works with init_db() pattern
     """
-    print("--- Running dedicated migration script ---")
-    # Assumes your DB file is named metadata.db in the root
+    print("--- Running migration script ---")
     db_path = str(ROOT_DIR / "metadata.db")
-    alembic_manager = AlembicMigrationManager(db_path)
+   
 
+    # Now handle the migration history
+    if handle_fresh_installation(db_path):
+        print("Fresh installation: database created and stamped successfully")
+        return
+    
+    # For existing installations with migration history, run normal migrations
+    print("Running incremental migrations...")
+    alembic_manager = AlembicMigrationManager(db_path)
     success, message = await alembic_manager.handle_database_upgrade()
 
     if not success:
         print(f"Migration Error: {message}")
-        # Exit with a non-zero status code to indicate failure
         sys.exit(1)
 
     print(f"Migration Success: {message}")
-    print("--- Migration script finished ---")
 
 if __name__ == "__main__":
     asyncio.run(main())