diff --git a/api/commands.py b/api/commands.py index 334e7daab5..8f11bc01d2 100644 --- a/api/commands.py +++ b/api/commands.py @@ -1,13 +1,13 @@ import base64 import json import logging +import os import secrets +import time +from pathlib import Path from typing import Optional import click -from flask import current_app -from werkzeug.exceptions import NotFound - from configs import dify_config from constants.languages import languages from core.rag.datasource.vdb.vector_factory import Vector @@ -16,15 +16,19 @@ 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 extensions.ext_storage import storage +from flask import current_app 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 from models import Tenant -from models.dataset import Dataset, DatasetCollectionBinding, DocumentSegment +from models.dataset import Dataset, DatasetCollectionBinding from models.dataset import Document as DatasetDocument -from models.model import Account, App, AppAnnotationSetting, AppMode, Conversation, MessageAnnotation +from models.dataset import DocumentSegment +from models.model import Account, App, AppAnnotationSetting, AppMode, Conversation, MessageAnnotation, UploadFile from models.provider import Provider, ProviderModel from services.account_service import RegisterService, TenantService +from werkzeug.exceptions import NotFound @click.command("reset-password", help="Reset the account password.") @@ -490,11 +494,10 @@ def add_qdrant_doc_id_index(field: str): click.echo(click.style("No dataset collection bindings found.", fg="red")) return import qdrant_client + from core.rag.datasource.vdb.qdrant.qdrant_vector import QdrantConfig from qdrant_client.http.exceptions import UnexpectedResponse from qdrant_client.http.models import PayloadSchemaType - from core.rag.datasource.vdb.qdrant.qdrant_vector import QdrantConfig - for binding in bindings: if dify_config.QDRANT_URL is None: raise ValueError("Qdrant URL is required.") @@ -507,7 +510,7 @@ def add_qdrant_doc_id_index(field: str): prefer_grpc=dify_config.QDRANT_GRPC_ENABLED, ) try: - client = qdrant_client.QdrantClient(**qdrant_config.to_qdrant_params()) + client = qdrant_client.QdrantClient(**qdrant_config.to_qdrant_params()) # type: ignore # create payload index client.create_payload_index(binding.collection_name, field, field_schema=PayloadSchemaType.KEYWORD) create_count += 1 @@ -649,3 +652,122 @@ where sites.id is null limit 1000""" break click.echo(click.style("Fix for missing app-related sites completed successfully!", fg="green")) + + +@click.command("upload-private-key-file-to-cloud-storage", help="upload private key file to cloud storage") +@click.option("--tenant_id", prompt=False, help="tenant_id") +def upload_private_key_file_cloud_storage(tenant_id: Optional[str] = None): + """ + upload private.pem to cloud storage + """ + click.echo( + click.style( + "Start upload private.pem to cloud storage", + fg="green", + ) + ) + + if not tenant_id: + click.echo( + click.style( + "Warning: did not provide an tenant_id, it will be auto queried in the database", + fg="yellow", + ) + ) + tenants_list: list[Tenant] = Tenant.query.all() + tenants = [item.id for item in tenants_list] + else: + tenants = [ + tenant_id, + ] + + for tenant_id in tenants: + click.echo( + click.style( + f"Current tenant_id is: {tenant_id}", + fg="green", + ) + ) + + file_key = f"privkeys/{tenant_id}/private.pem" + file_content = Path(f"{os.environ.get('STORAGE_LOCAL_PATH', 'storage')}/{file_key}").read_bytes() + storage.save(filename=file_key, data=file_content) + click.echo( + click.style( + f"Congratulations! file uploaded. file.key: {file_key}", + fg="green", + ) + ) + + +@click.command("upload-local-files-to-cloud-storage", help="upload local files to cloud storage") +def upload_local_files_to_cloud_storage(): + """ + upload local files to cloud storage + """ + click.echo( + click.style( + "Start upload local files to cloud storage", + fg="green", + ) + ) + + total_count = UploadFile.query.filter_by(storage_type="local").count() + click.echo(click.style(f"Total files to process: {total_count}", fg="green")) + + batch_size = 100 + processed_count = 0 + while processed_count < total_count: + files: list[UploadFile] = UploadFile.query.filter_by(storage_type="local").limit(batch_size).all() + + for file in files: + target_filepath = f"{os.environ.get('STORAGE_LOCAL_PATH', 'storage')}/{file.key}" + + # if the file exists + if not os.path.exists(target_filepath): + click.echo( + click.style( + f"Warning! file not exist. filepath: {target_filepath}, ignore this, continue", + fg="yellow", + ) + ) + processed_count += 1 + if processed_count % 10 == 0 or processed_count == total_count: + click.echo(click.style(f"Processed {processed_count}/{total_count} files\n", fg="blue")) + continue + + # Upload to cloud storage + file_content = Path(target_filepath).read_bytes() + storage.save(filename=file.key, data=file_content) + click.echo( + click.style( + f"File uploaded. file.key: {file.key}", + fg="green", + ) + ) + + # Update database record + try: + file.storage_type = os.environ["STORAGE_TYPE"] + db.session.commit() + click.echo( + click.style( + f"file.storage_type updated to database. file.key: {file.key}", + fg="green", + ) + ) + except Exception as e: + click.echo(click.style(f"An error occurred: {str(e)}", fg="red")) + db.session.rollback() + + processed_count += 1 + if processed_count % 10 == 0 or processed_count == total_count: + click.echo(click.style(f"Processed {processed_count}/{total_count} files\n", fg="blue")) + + time.sleep(3) + click.echo( + click.style( + "Congratulations! finish files uploaded.", + fg="green", + ) + ) diff --git a/api/extensions/ext_commands.py b/api/extensions/ext_commands.py index ccf0d316ca..d96b3e3b49 100644 --- a/api/extensions/ext_commands.py +++ b/api/extensions/ext_commands.py @@ -11,6 +11,8 @@ def init_app(app: DifyApp): reset_encrypt_key_pair, reset_password, upgrade_db, + upload_local_files_to_cloud_storage, + upload_private_key_file_cloud_storage, vdb_migrate, ) @@ -24,6 +26,8 @@ def init_app(app: DifyApp): create_tenant, upgrade_db, fix_app_site_missing, + upload_private_key_file_cloud_storage, + upload_local_files_to_cloud_storage, ] for cmd in cmds_to_register: app.cli.add_command(cmd)