m

.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/[email protected]
         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 db-migrate

api/app.py

@@ -267,5 +267,9 @@ def pool_stat():
     }
 
 
+def db_migrate():
+    with app.app_context():
+        ext_migrate.db_upgrade()
+
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=5001)

api/commands.py

@@ -12,6 +12,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 +554,22 @@ def create_tenant(email: str, language: Optional[str] = None):
                            'Account: {}\nPassword: {}'.format(email, new_password), fg='green'))
 
 
+@click.command('db-migrate', help='migrate the database')
+def db_migrate():
+    click.echo('Preparing database migration...')
+    lock = redis_client.lock(name='db_upgrade_lock', timeout=60)
+    if lock.acquire(blocking=False):
+        try:
+            import flask_migrate
+            click.echo(click.style('Start database migration.', fg='green'))
+            flask_migrate.upgrade()
+            click.echo(click.style('Database migration successful!', fg='green'))
+        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 +578,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(db_migrate)

api/docker/entrypoint.sh

@@ -4,7 +4,7 @@ set -e
 
 if [[ "${MIGRATION_ENABLED}" == "true" ]]; then
   echo "Running migrations"
-  flask db upgrade
+  flask db-migrate
 fi
 
 if [[ "${MODE}" == "worker" ]]; then

api/extensions/ext_migrate.py

@@ -3,3 +3,10 @@
 
 def init(app, db):
     flask_migrate.Migrate(app, db)
+
+
+def db_upgrade():
+    """
+    Upgrade the database with Flask-Migrate
+    """
+    flask_migrate.upgrade()

api/migrations/README

@@ -1,2 +1 @@
 Single-database configuration for Flask.
-

api/pyproject.toml

@@ -91,6 +91,9 @@ authors = ["Dify <[email protected]>"]
 readme = "README.md"
 package-mode = false
 
+[tool.poetry.scripts]
+db_migration = 'app:db_migration'
+
 [tool.poetry.dependencies]
 python = "^3.10"
 beautifulsoup4 = "4.12.2"