diff --git a/dbgpt/app/_cli.py b/dbgpt/app/_cli.py index 98bed72cf..f48e57111 100644 --- a/dbgpt/app/_cli.py +++ b/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 diff --git a/dbgpt/app/initialization/db_model_initialization.py b/dbgpt/app/initialization/db_model_initialization.py index 831b76fec..9095da3e0 100644 --- a/dbgpt/app/initialization/db_model_initialization.py +++ b/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, diff --git a/dbgpt/util/_db_migration_utils.py b/dbgpt/util/_db_migration_utils.py index 987606364..13734960d 100644 --- a/dbgpt/util/_db_migration_utils.py +++ b/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: