Repository: langgenius/dify
File: api/commands.py (Python)

Code before:
        password=new_password,
        language=language,
        create_workspace_required=False,
    )
    TenantService.create_owner_tenant_if_not_exist(account, name)

    click.echo(
        click.style(
            f"Account and tenant created.\nAccount: {email}\nPassword: {new_password}",
            fg="green",
        )
    )


@click.command("upgrade-db", help="Upgrade the database")
def upgrade_db():
    click.echo("Preparing database migration...")
    # Use a short base TTL + heartbeat renewal, so a crashed process doesn't block migrations for long.
    # thread_local=False is required because heartbeat runs in a separate thread.
    lock = redis_client.lock(
        name="db_upgrade_lock",
        timeout=DB_UPGRADE_LOCK_TTL_SECONDS,
        thread_local=False,
    )
    if lock.acquire(blocking=False):
        stop_event = threading.Event()
        heartbeat_thread = threading.Thread(
            target=_heartbeat_db_upgrade_lock,
            args=(lock, stop_event, float(DB_UPGRADE_LOCK_TTL_SECONDS)),
            daemon=True,
        )
        heartbeat_thread.start()
        migration_succeeded = False
        try:
            click.echo(click.style("Starting database migration.", fg="green"))

            # run db migration
            import flask_migrate

            flask_migrate.upgrade()

            migration_succeeded = True
            click.echo(click.style("Database migration successful!", fg="green"))

        except Exception as e:
            logger.exception("Failed to execute database migration")
            click.echo(click.style(f"Database migration failed: {e}", fg="red"))
            raise SystemExit(1)
        finally:
            stop_event.set()
            heartbeat_thread.join(timeout=5)

Review comment (suggestion):
Consider packing it into a general clean API, so that invoker don't need to consider too much on the threading things, for example:

```suggestion
lock = AutoRenewRedisLock(name="xxx", ttl=xxx, ...)
lock.acquire(blocking=False)
```

Code after:

            # run db migration
            import flask_migrate

            flask_migrate.upgrade()

            migration_succeeded = True
            click.echo(click.style("Database migration successful!", fg="green"))

        except Exception as e:
            logger.exception("Failed to execute database migration")
            click.echo(click.style(f"Database migration failed: {e}", fg="red"))
            raise SystemExit(1)
        finally:
            status = "successful" if migration_succeeded else "failed"
            lock.release_safely(status=status)
    else:
        click.echo("Database migration skipped")


@click.command("fix-app-site-missing", help="Fix app related site missing issue.")
def fix_app_site_missing():
    """
    Fix app related site missing issue.
    """
    click.echo(click.style("Starting fix for missing app-related sites.", fg="green"))

    failed_app_ids = []
    while True:
        sql = """select apps.id as id from apps left join sites on sites.app_id=apps.id
where sites.id is null limit 1000"""
        with db.engine.begin() as conn:
            rs = conn.execute(sa.text(sql))

            processed_count = 0
            for i in rs:
                processed_count += 1
                app_id = str(i.id)

                if app_id in failed_app_ids:
                    continue

                try:
                    app = db.session.query(App).where(App.id == app_id).first()
                    if not app:
                        logger.info("App %s not found", app_id)
                        continue

                    tenant = app.tenant
                    if tenant:
                        accounts = tenant.get_accounts()