Merge branch 'main' into fix/chore-fix

api/tasks/add_document_to_index_task.py

@@ -38,7 +38,11 @@ def add_document_to_index_task(dataset_document_id: str):
     try:
         segments = (
             db.session.query(DocumentSegment)
-            .filter(DocumentSegment.document_id == dataset_document.id, DocumentSegment.enabled == True)
+            .filter(
+                DocumentSegment.document_id == dataset_document.id,
+                DocumentSegment.enabled == False,
+                DocumentSegment.status == "completed",
+            )
             .order_by(DocumentSegment.position.asc())
             .all()
         )
@@ -85,6 +89,16 @@ def add_document_to_index_task(dataset_document_id: str):
         db.session.query(DatasetAutoDisableLog).filter(
             DatasetAutoDisableLog.document_id == dataset_document.id
         ).delete()
+
+        # update segment to enable
+        db.session.query(DocumentSegment).filter(DocumentSegment.document_id == dataset_document.id).update(
+            {
+                DocumentSegment.enabled: True,
+                DocumentSegment.disabled_at: None,
+                DocumentSegment.disabled_by: None,
+                DocumentSegment.updated_at: datetime.datetime.now(datetime.UTC).replace(tzinfo=None),
+            }
+        )
         db.session.commit()
 
         end_at = time.perf_counter()

api/tasks/delete_account_task.py

@@ -0,0 +1,26 @@
+import logging
+
+from celery import shared_task  # type: ignore
+
+from extensions.ext_database import db
+from models.account import Account
+from services.billing_service import BillingService
+from tasks.mail_account_deletion_task import send_deletion_success_task
+
+logger = logging.getLogger(__name__)
+
+
+@shared_task(queue="dataset")
+def delete_account_task(account_id):
+    account = db.session.query(Account).filter(Account.id == account_id).first()
+    try:
+        BillingService.delete_account(account_id)
+    except Exception as e:
+        logger.exception(f"Failed to delete account {account_id} from billing service.")
+        raise
+
+    if not account:
+        logger.error(f"Account {account_id} not found.")
+        return
+    # send success email
+    send_deletion_success_task.delay(account.email)

api/tasks/mail_account_deletion_task.py

@@ -0,0 +1,70 @@
+import logging
+import time
+
+import click
+from celery import shared_task  # type: ignore
+from flask import render_template
+
+from extensions.ext_mail import mail
+
+
+@shared_task(queue="mail")
+def send_deletion_success_task(to):
+    """Send email to user regarding account deletion.
+
+    Args:
+        log (AccountDeletionLog): Account deletion log object
+    """
+    if not mail.is_inited():
+        return
+
+    logging.info(click.style(f"Start send account deletion success email to {to}", fg="green"))
+    start_at = time.perf_counter()
+
+    try:
+        html_content = render_template(
+            "delete_account_success_template_en-US.html",
+            to=to,
+            email=to,
+        )
+        mail.send(to=to, subject="Your Dify.AI Account Has Been Successfully Deleted", html=html_content)
+
+        end_at = time.perf_counter()
+        logging.info(
+            click.style(
+                "Send account deletion success email to {}: latency: {}".format(to, end_at - start_at), fg="green"
+            )
+        )
+    except Exception:
+        logging.exception("Send account deletion success email to {} failed".format(to))
+
+
+@shared_task(queue="mail")
+def send_account_deletion_verification_code(to, code):
+    """Send email to user regarding account deletion verification code.
+
+    Args:
+        to (str): Recipient email address
+        code (str): Verification code
+    """
+    if not mail.is_inited():
+        return
+
+    logging.info(click.style(f"Start send account deletion verification code email to {to}", fg="green"))
+    start_at = time.perf_counter()
+
+    try:
+        html_content = render_template("delete_account_code_email_template_en-US.html", to=to, code=code)
+        mail.send(to=to, subject="Dify.AI Account Deletion and Verification", html=html_content)
+
+        end_at = time.perf_counter()
+        logging.info(
+            click.style(
+                "Send account deletion verification code email to {} succeeded: latency: {}".format(
+                    to, end_at - start_at
+                ),
+                fg="green",
+            )
+        )
+    except Exception:
+        logging.exception("Send account deletion verification code email to {} failed".format(to))

api/tasks/remove_document_from_index_task.py

@@ -1,3 +1,4 @@
+import datetime
 import logging
 import time
 
@@ -46,6 +47,16 @@ def remove_document_from_index_task(document_id: str):
                 index_processor.clean(dataset, index_node_ids, with_keywords=True, delete_child_chunks=False)
             except Exception:
                 logging.exception(f"clean dataset {dataset.id} from index failed")
+        # update segment to disable
+        db.session.query(DocumentSegment).filter(DocumentSegment.document_id == document.id).update(
+            {
+                DocumentSegment.enabled: False,
+                DocumentSegment.disabled_at: datetime.datetime.now(datetime.UTC).replace(tzinfo=None),
+                DocumentSegment.disabled_by: document.disabled_by,
+                DocumentSegment.updated_at: datetime.datetime.now(datetime.UTC).replace(tzinfo=None),
+            }
+        )
+        db.session.commit()
 
         end_at = time.perf_counter()
         logging.info(