feat: add fix-app-site-missing command (#5711)

This commit is contained in:
takatost 2024-06-28 20:20:23 +08:00 committed by GitHub
parent a2c260fba0
commit d30c13891b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 43 additions and 2 deletions

View File

@ -12,6 +12,7 @@ from constants.languages import languages
from core.rag.datasource.vdb.vector_factory import Vector
from core.rag.datasource.vdb.vector_type import VectorType
from core.rag.models.document import Document
from events.app_event import app_was_created
from extensions.ext_database import db
from extensions.ext_redis import redis_client
from libs.helper import email as email_validate
@ -585,6 +586,46 @@ def upgrade_db():
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('Start fix app related site missing issue.', fg='green'))
while True:
try:
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(db.text(sql))
processed_count = 0
for i in rs:
processed_count += 1
app_id = str(i.id)
app = db.session.query(App).filter(App.id == app_id).first()
tenant = app.tenant
if tenant:
accounts = tenant.get_accounts()
if not accounts:
print("Fix app {} failed.".format(app.id))
continue
account = accounts[0]
print("Fix app {} related site missing issue.".format(app.id))
app_was_created.send(app, account=account)
if not processed_count:
break
except Exception as e:
click.echo(click.style('Fix app related site missing issue failed!', fg='red'))
logging.exception(f'Fix app related site missing issue failed, error: {e}')
continue
click.echo(click.style('Congratulations! Fix app related site missing issue successful!', fg='green'))
def register_commands(app):
app.cli.add_command(reset_password)
app.cli.add_command(reset_email)
@ -594,3 +635,4 @@ def register_commands(app):
app.cli.add_command(add_qdrant_doc_id_index)
app.cli.add_command(create_tenant)
app.cli.add_command(upgrade_db)
app.cli.add_command(fix_app_site_missing)

View File

@ -153,8 +153,7 @@ class Tenant(db.Model):
created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)'))
updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)'))
def get_accounts(self) -> list[db.Model]:
Account = db.Model
def get_accounts(self) -> list[Account]:
return db.session.query(Account).filter(
Account.id == TenantAccountJoin.account_id,
TenantAccountJoin.tenant_id == self.id