From 8fad27d904eb39ed0f29c76419701aa1850e93ea Mon Sep 17 00:00:00 2001
From: YueGuobin <yueguobin@outlook.com>
Date: Tue, 3 Mar 2026 16:40:57 +0800
Subject: [PATCH] feat(db): improve database initialization and migration
 handling

- Add smarter database state detection during initialization to handle new databases, existing databases with new features, and old databases needing migration
- Make migration for llm_model_configs table idempotent to prevent conflicts when table already exists from code
- Add detailed logging for different database initialization scenarios
- Import sqlalchemy module for database inspection capabilities
---
 gns3server/db/tasks.py                        | 44 ++++++++++++++++++-
 ...20260303_create_llm_model_configs_table.py | 12 +++++
 2 files changed, 54 insertions(+), 2 deletions(-)

diff --git a/gns3server/db/tasks.py b/gns3server/db/tasks.py
index bfb7dac48..b2458a219 100644
--- a/gns3server/db/tasks.py
+++ b/gns3server/db/tasks.py
@@ -25,6 +25,7 @@ from typing import List
 from sqlalchemy import event
 from sqlalchemy.engine import Engine
 from sqlalchemy.exc import SQLAlchemyError
+import sqlalchemy as sa
 from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
 from alembic import command, config
 from alembic.script import ScriptDirectory
@@ -87,12 +88,51 @@ async def connect_to_db(app: FastAPI) -> None:
             current_rev, head_rev = await conn.run_sync(check_revision, alembic_cfg)
             log.info(f"Current database revision is {current_rev}")
             if current_rev is None:
-                await conn.run_sync(Base.metadata.create_all)
-                await conn.run_sync(run_stamp, alembic_cfg)
+                # No version tracking found. Check if this is a truly new database
+                # or an old database that needs migration.
+                def check_db_state(connection):
+                    # Check if llm_model_configs table exists
+                    inspector = sa.inspect(connection)
+                    tables = inspector.get_table_names()
+
+                    if 'users' not in tables:
+                        return 'new'  # Truly new database
+
+                    # Check for new feature columns that indicate this is already migrated
+                    columns = [col['name'] for col in inspector.get_columns('users')]
+                    if 'llm_model_configs' in tables:
+                        # The llm_model_configs table already exists (created from code)
+                        return 'new_with_llm_configs'
+                    else:
+                        # Old database without llm_model_configs table, needs migration
+                        return 'old_needs_migration'
+
+                db_state = await conn.run_sync(check_db_state)
+
+                if db_state == 'new':
+                    # Truly new database: create all tables and stamp
+                    await conn.run_sync(Base.metadata.create_all)
+                    await conn.run_sync(run_stamp, alembic_cfg)
+                    await conn.commit()
+                    log.info("Created new database and stamped to head revision")
+                elif db_state == 'new_with_llm_configs':
+                    # Database already has llm_model_configs table (from Base.metadata.create_all)
+                    # Just stamp the version
+                    await conn.run_sync(run_stamp, alembic_cfg)
+                    await conn.commit()
+                    log.info("Database has llm_model_configs table, stamped to head revision")
+                else:
+                    # Old database without llm_model_configs table: run migrations
+                    log.info("Old database detected, running migrations to add llm_model_configs table...")
+                    await conn.run_sync(run_upgrade, alembic_cfg)
+                    await conn.commit()
+                    log.info("Database migrations completed successfully")
             elif current_rev != head_rev:
                 # upgrade the database if needed
+                log.info(f"Upgrading database from revision {current_rev} to {head_rev}...")
                 await conn.run_sync(run_upgrade, alembic_cfg)
                 await conn.commit()
+                log.info("Database upgrade completed successfully")
         app.state._db_engine = engine
     except SQLAlchemyError as e:
         log.fatal(f"Error while connecting to database '{db_url}: {e}")
diff --git a/gns3server/db_migrations/versions/20260303_create_llm_model_configs_table.py b/gns3server/db_migrations/versions/20260303_create_llm_model_configs_table.py
index 45a19ab9a..330e06f39 100644
--- a/gns3server/db_migrations/versions/20260303_create_llm_model_configs_table.py
+++ b/gns3server/db_migrations/versions/20260303_create_llm_model_configs_table.py
@@ -8,6 +8,7 @@ Create Date: 2026-03-03
 from alembic import op
 import sqlalchemy as sa
 from sqlalchemy.dialects import postgresql
+from sqlalchemy import inspect
 
 # revision identifiers, used by Alembic.
 revision = '20260303_create_llm_model_configs'
@@ -17,6 +18,17 @@ depends_on = None
 
 
 def upgrade() -> None:
+    # Get the current connection
+    conn = op.get_bind()
+    inspector = inspect(conn)
+
+    # Check if table already exists (idempotent for databases created from code)
+    tables = inspector.get_table_names()
+
+    if 'llm_model_configs' in tables:
+        # Table already exists from Base.metadata.create_all, skip creation
+        return
+
     # Create llm_model_configs table
     op.create_table(
         'llm_model_configs',