From d30c13891b37608003e920b7ce07a3fbfa1f35cc Mon Sep 17 00:00:00 2001
From: takatost <takatost@users.noreply.github.com>
Date: Fri, 28 Jun 2024 20:20:23 +0800
Subject: [PATCH] feat: add fix-app-site-missing command (#5711)

---
 api/commands.py       | 42 ++++++++++++++++++++++++++++++++++++++++++
 api/models/account.py |  3 +--
 2 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/api/commands.py b/api/commands.py
index 91d7737023..fa6221df52 100644
--- a/api/commands.py
+++ b/api/commands.py
@@ -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)
diff --git a/api/models/account.py b/api/models/account.py
index 4911757b07..3b258c4c82 100644
--- a/api/models/account.py
+++ b/api/models/account.py
@@ -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