diff --git a/.github/workflows/db-migration-test.yml b/.github/workflows/db-migration-test.yml index 2f059bea9d7aae..64e8eb291c466e 100644 --- a/.github/workflows/db-migration-test.yml +++ b/.github/workflows/db-migration-test.yml @@ -38,13 +38,14 @@ jobs: - name: Install dependencies run: poetry install -C api - - name: Set up Middleware + - name: Set up Middlewares uses: hoverkraft-tech/compose-action@v2.0.0 with: compose-file: | docker/docker-compose.middleware.yaml services: | db + redis - name: Prepare configs run: | @@ -54,4 +55,4 @@ jobs: - name: Run DB Migration run: | cd api - poetry run python -m flask db upgrade + poetry run python -m flask upgrade-db diff --git a/api/commands.py b/api/commands.py index 7a8d8000a0f443..f3e0769134f137 100644 --- a/api/commands.py +++ b/api/commands.py @@ -1,5 +1,6 @@ import base64 import json +import logging import secrets from typing import Optional @@ -12,6 +13,7 @@ from core.rag.datasource.vdb.vector_type import VectorType from core.rag.models.document import Document from extensions.ext_database import db +from extensions.ext_redis import redis_client from libs.helper import email as email_validate from libs.password import hash_password, password_pattern, valid_password from libs.rsa import generate_key_pair @@ -553,6 +555,28 @@ def create_tenant(email: str, language: Optional[str] = None): 'Account: {}\nPassword: {}'.format(email, new_password), fg='green')) +@click.command('upgrade-db', help='upgrade the database') +def upgrade_db(): + click.echo('Preparing database migration...') + lock = redis_client.lock(name='db_upgrade_lock', timeout=60) + if lock.acquire(blocking=False): + try: + click.echo(click.style('Start database migration.', fg='green')) + + # run db migration + import flask_migrate + flask_migrate.upgrade() + + click.echo(click.style('Database migration successful!', fg='green')) + + except Exception as e: + logging.exception(f'Database migration failed, error: {e}') + finally: + lock.release() + else: + click.echo('Database migration skipped') + + def register_commands(app): app.cli.add_command(reset_password) app.cli.add_command(reset_email) @@ -561,4 +585,4 @@ def register_commands(app): app.cli.add_command(convert_to_agent_apps) app.cli.add_command(add_qdrant_doc_id_index) app.cli.add_command(create_tenant) - + app.cli.add_command(upgrade_db) diff --git a/api/docker/entrypoint.sh b/api/docker/entrypoint.sh index 2dad52f8f6d2bf..7b1f1dfc03c679 100755 --- a/api/docker/entrypoint.sh +++ b/api/docker/entrypoint.sh @@ -4,7 +4,7 @@ set -e if [[ "${MIGRATION_ENABLED}" == "true" ]]; then echo "Running migrations" - flask db upgrade + flask upgrade-db fi if [[ "${MODE}" == "worker" ]]; then diff --git a/api/migrations/README b/api/migrations/README index 220678df7ab06e..0e048441597444 100644 --- a/api/migrations/README +++ b/api/migrations/README @@ -1,2 +1 @@ Single-database configuration for Flask. -