add oss cmd

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",
+        )
+    )

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)