fix: Fix migration bug

dbgpt/app/_cli.py

@@ -272,8 +272,8 @@ def _get_migration_config(
     from dbgpt.storage.metadata.db_manager import db as db_manager
     from dbgpt.util._db_migration_utils import create_alembic_config
 
-    # Must import dbgpt_server for initialize db metadata
-    from dbgpt.app.dbgpt_server import initialize_app as _
+    # Import all models to make sure they are registered with SQLAlchemy.
+    from dbgpt.app.initialization.db_model_initialization import _MODELS
     from dbgpt.app.base import _initialize_db
 
     # initialize db

dbgpt/app/initialization/db_model_initialization.py

@@ -8,6 +8,7 @@
 from dbgpt.app.openapi.api_v1.feedback.feed_back_db import ChatFeedBackEntity
 
 # from dbgpt.app.prompt.prompt_manage_db import PromptManageEntity
+from dbgpt.serve.prompt.models.models import ServeEntity as PromptManageEntity
 from dbgpt.datasource.manages.connect_config_db import ConnectConfigEntity
 from dbgpt.storage.chat_history.chat_history_db import (
     ChatHistoryEntity,
@@ -17,7 +18,7 @@
 _MODELS = [
     PluginHubEntity,
     MyPluginEntity,
-    # PromptManageEntity,
+    PromptManageEntity,
     KnowledgeSpaceEntity,
     KnowledgeDocumentEntity,
     DocumentChunkEntity,

dbgpt/util/_db_migration_utils.py

@@ -77,8 +77,13 @@ def create_migration_script(
         current_rev = context.get_current_revision()
         head_rev = script_dir.get_current_head()
 
+    logger.info(
+        f"alembic migration current revision: {current_rev}, latest revision: {head_rev}"
+    )
     should_create_revision = (
-        current_rev != head_rev or create_new_revision_if_noting_to_update
+        (current_rev is None and head_rev is None)
+        or current_rev != head_rev
+        or create_new_revision_if_noting_to_update
     )
     if should_create_revision:
         with engine.connect() as connection:
@@ -265,6 +270,7 @@ def get_current_revision(engine):
     if current_rev != head_rev:
         logger.error(
             "Database is not at the latest revision. "
+            f"Current revision: {current_rev}, latest revision: {head_rev}\n"
             "Please apply existing migration scripts before generating new ones. "
             "Check the listed file paths for migration scripts.\n"
             f"Also you can try the following solutions:\n{_MIGRATION_SOLUTION}\n"
@@ -353,8 +359,9 @@ def _ddl_init_and_upgrade(
     try:
         latest_revision_before = _get_latest_revision(alembic_cfg, db.engine)
         # create_new_revision_if_noting_to_update=False avoid creating a lot of empty migration scripts
+        # TODO Set create_new_revision_if_noting_to_update=False, not working now.
         new_script_path = create_migration_script(
-            alembic_cfg, db.engine, create_new_revision_if_noting_to_update=False
+            alembic_cfg, db.engine, create_new_revision_if_noting_to_update=True
         )
         upgrade_database(alembic_cfg, db.engine)
     except CommandError as e: