commands.py 31 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794
  1. import base64
  2. import json
  3. import logging
  4. import secrets
  5. from typing import Optional
  6. import click
  7. from flask import current_app
  8. from werkzeug.exceptions import NotFound
  9. from configs import dify_config
  10. from constants.languages import languages
  11. from core.rag.datasource.vdb.vector_factory import Vector
  12. from core.rag.datasource.vdb.vector_type import VectorType
  13. from core.rag.index_processor.constant.built_in_field import BuiltInField
  14. from core.rag.models.document import Document
  15. from events.app_event import app_was_created
  16. from extensions.ext_database import db
  17. from extensions.ext_redis import redis_client
  18. from libs.helper import email as email_validate
  19. from libs.password import hash_password, password_pattern, valid_password
  20. from libs.rsa import generate_key_pair
  21. from models import Tenant
  22. from models.dataset import Dataset, DatasetCollectionBinding, DatasetMetadata, DatasetMetadataBinding, DocumentSegment
  23. from models.dataset import Document as DatasetDocument
  24. from models.model import Account, App, AppAnnotationSetting, AppMode, Conversation, MessageAnnotation
  25. from models.provider import Provider, ProviderModel
  26. from services.account_service import RegisterService, TenantService
  27. from services.plugin.data_migration import PluginDataMigration
  28. from services.plugin.plugin_migration import PluginMigration
  29. @click.command("reset-password", help="Reset the account password.")
  30. @click.option("--email", prompt=True, help="Account email to reset password for")
  31. @click.option("--new-password", prompt=True, help="New password")
  32. @click.option("--password-confirm", prompt=True, help="Confirm new password")
  33. def reset_password(email, new_password, password_confirm):
  34. """
  35. Reset password of owner account
  36. Only available in SELF_HOSTED mode
  37. """
  38. if str(new_password).strip() != str(password_confirm).strip():
  39. click.echo(click.style("Passwords do not match.", fg="red"))
  40. return
  41. account = db.session.query(Account).filter(Account.email == email).one_or_none()
  42. if not account:
  43. click.echo(click.style("Account not found for email: {}".format(email), fg="red"))
  44. return
  45. try:
  46. valid_password(new_password)
  47. except:
  48. click.echo(click.style("Invalid password. Must match {}".format(password_pattern), fg="red"))
  49. return
  50. # generate password salt
  51. salt = secrets.token_bytes(16)
  52. base64_salt = base64.b64encode(salt).decode()
  53. # encrypt password with salt
  54. password_hashed = hash_password(new_password, salt)
  55. base64_password_hashed = base64.b64encode(password_hashed).decode()
  56. account.password = base64_password_hashed
  57. account.password_salt = base64_salt
  58. db.session.commit()
  59. click.echo(click.style("Password reset successfully.", fg="green"))
  60. @click.command("reset-email", help="Reset the account email.")
  61. @click.option("--email", prompt=True, help="Current account email")
  62. @click.option("--new-email", prompt=True, help="New email")
  63. @click.option("--email-confirm", prompt=True, help="Confirm new email")
  64. def reset_email(email, new_email, email_confirm):
  65. """
  66. Replace account email
  67. :return:
  68. """
  69. if str(new_email).strip() != str(email_confirm).strip():
  70. click.echo(click.style("New emails do not match.", fg="red"))
  71. return
  72. account = db.session.query(Account).filter(Account.email == email).one_or_none()
  73. if not account:
  74. click.echo(click.style("Account not found for email: {}".format(email), fg="red"))
  75. return
  76. try:
  77. email_validate(new_email)
  78. except:
  79. click.echo(click.style("Invalid email: {}".format(new_email), fg="red"))
  80. return
  81. account.email = new_email
  82. db.session.commit()
  83. click.echo(click.style("Email updated successfully.", fg="green"))
  84. @click.command(
  85. "reset-encrypt-key-pair",
  86. help="Reset the asymmetric key pair of workspace for encrypt LLM credentials. "
  87. "After the reset, all LLM credentials will become invalid, "
  88. "requiring re-entry."
  89. "Only support SELF_HOSTED mode.",
  90. )
  91. @click.confirmation_option(
  92. prompt=click.style(
  93. "Are you sure you want to reset encrypt key pair? This operation cannot be rolled back!", fg="red"
  94. )
  95. )
  96. def reset_encrypt_key_pair():
  97. """
  98. Reset the encrypted key pair of workspace for encrypt LLM credentials.
  99. After the reset, all LLM credentials will become invalid, requiring re-entry.
  100. Only support SELF_HOSTED mode.
  101. """
  102. if dify_config.EDITION != "SELF_HOSTED":
  103. click.echo(click.style("This command is only for SELF_HOSTED installations.", fg="red"))
  104. return
  105. tenants = db.session.query(Tenant).all()
  106. for tenant in tenants:
  107. if not tenant:
  108. click.echo(click.style("No workspaces found. Run /install first.", fg="red"))
  109. return
  110. tenant.encrypt_public_key = generate_key_pair(tenant.id)
  111. db.session.query(Provider).filter(Provider.provider_type == "custom", Provider.tenant_id == tenant.id).delete()
  112. db.session.query(ProviderModel).filter(ProviderModel.tenant_id == tenant.id).delete()
  113. db.session.commit()
  114. click.echo(
  115. click.style(
  116. "Congratulations! The asymmetric key pair of workspace {} has been reset.".format(tenant.id),
  117. fg="green",
  118. )
  119. )
  120. @click.command("vdb-migrate", help="Migrate vector db.")
  121. @click.option("--scope", default="all", prompt=False, help="The scope of vector database to migrate, Default is All.")
  122. def vdb_migrate(scope: str):
  123. if scope in {"knowledge", "all"}:
  124. migrate_knowledge_vector_database()
  125. if scope in {"annotation", "all"}:
  126. migrate_annotation_vector_database()
  127. def migrate_annotation_vector_database():
  128. """
  129. Migrate annotation datas to target vector database .
  130. """
  131. click.echo(click.style("Starting annotation data migration.", fg="green"))
  132. create_count = 0
  133. skipped_count = 0
  134. total_count = 0
  135. page = 1
  136. while True:
  137. try:
  138. # get apps info
  139. per_page = 50
  140. apps = (
  141. db.session.query(App)
  142. .filter(App.status == "normal")
  143. .order_by(App.created_at.desc())
  144. .limit(per_page)
  145. .offset((page - 1) * per_page)
  146. .all()
  147. )
  148. if not apps:
  149. break
  150. except NotFound:
  151. break
  152. page += 1
  153. for app in apps:
  154. total_count = total_count + 1
  155. click.echo(
  156. f"Processing the {total_count} app {app.id}. " + f"{create_count} created, {skipped_count} skipped."
  157. )
  158. try:
  159. click.echo("Creating app annotation index: {}".format(app.id))
  160. app_annotation_setting = (
  161. db.session.query(AppAnnotationSetting).filter(AppAnnotationSetting.app_id == app.id).first()
  162. )
  163. if not app_annotation_setting:
  164. skipped_count = skipped_count + 1
  165. click.echo("App annotation setting disabled: {}".format(app.id))
  166. continue
  167. # get dataset_collection_binding info
  168. dataset_collection_binding = (
  169. db.session.query(DatasetCollectionBinding)
  170. .filter(DatasetCollectionBinding.id == app_annotation_setting.collection_binding_id)
  171. .first()
  172. )
  173. if not dataset_collection_binding:
  174. click.echo("App annotation collection binding not found: {}".format(app.id))
  175. continue
  176. annotations = db.session.query(MessageAnnotation).filter(MessageAnnotation.app_id == app.id).all()
  177. dataset = Dataset(
  178. id=app.id,
  179. tenant_id=app.tenant_id,
  180. indexing_technique="high_quality",
  181. embedding_model_provider=dataset_collection_binding.provider_name,
  182. embedding_model=dataset_collection_binding.model_name,
  183. collection_binding_id=dataset_collection_binding.id,
  184. )
  185. documents = []
  186. if annotations:
  187. for annotation in annotations:
  188. document = Document(
  189. page_content=annotation.question,
  190. metadata={"annotation_id": annotation.id, "app_id": app.id, "doc_id": annotation.id},
  191. )
  192. documents.append(document)
  193. vector = Vector(dataset, attributes=["doc_id", "annotation_id", "app_id"])
  194. click.echo(f"Migrating annotations for app: {app.id}.")
  195. try:
  196. vector.delete()
  197. click.echo(click.style(f"Deleted vector index for app {app.id}.", fg="green"))
  198. except Exception as e:
  199. click.echo(click.style(f"Failed to delete vector index for app {app.id}.", fg="red"))
  200. raise e
  201. if documents:
  202. try:
  203. click.echo(
  204. click.style(
  205. f"Creating vector index with {len(documents)} annotations for app {app.id}.",
  206. fg="green",
  207. )
  208. )
  209. vector.create(documents)
  210. click.echo(click.style(f"Created vector index for app {app.id}.", fg="green"))
  211. except Exception as e:
  212. click.echo(click.style(f"Failed to created vector index for app {app.id}.", fg="red"))
  213. raise e
  214. click.echo(f"Successfully migrated app annotation {app.id}.")
  215. create_count += 1
  216. except Exception as e:
  217. click.echo(
  218. click.style(
  219. "Error creating app annotation index: {} {}".format(e.__class__.__name__, str(e)), fg="red"
  220. )
  221. )
  222. continue
  223. click.echo(
  224. click.style(
  225. f"Migration complete. Created {create_count} app annotation indexes. Skipped {skipped_count} apps.",
  226. fg="green",
  227. )
  228. )
  229. def migrate_knowledge_vector_database():
  230. """
  231. Migrate vector database datas to target vector database .
  232. """
  233. click.echo(click.style("Starting vector database migration.", fg="green"))
  234. create_count = 0
  235. skipped_count = 0
  236. total_count = 0
  237. vector_type = dify_config.VECTOR_STORE
  238. upper_collection_vector_types = {
  239. VectorType.MILVUS,
  240. VectorType.PGVECTOR,
  241. VectorType.RELYT,
  242. VectorType.WEAVIATE,
  243. VectorType.ORACLE,
  244. VectorType.ELASTICSEARCH,
  245. VectorType.OPENGAUSS,
  246. }
  247. lower_collection_vector_types = {
  248. VectorType.ANALYTICDB,
  249. VectorType.CHROMA,
  250. VectorType.MYSCALE,
  251. VectorType.PGVECTO_RS,
  252. VectorType.TIDB_VECTOR,
  253. VectorType.OPENSEARCH,
  254. VectorType.TENCENT,
  255. VectorType.BAIDU,
  256. VectorType.VIKINGDB,
  257. VectorType.UPSTASH,
  258. VectorType.COUCHBASE,
  259. VectorType.OCEANBASE,
  260. }
  261. page = 1
  262. while True:
  263. try:
  264. datasets = (
  265. Dataset.query.filter(Dataset.indexing_technique == "high_quality")
  266. .order_by(Dataset.created_at.desc())
  267. .paginate(page=page, per_page=50)
  268. )
  269. except NotFound:
  270. break
  271. page += 1
  272. for dataset in datasets:
  273. total_count = total_count + 1
  274. click.echo(
  275. f"Processing the {total_count} dataset {dataset.id}. {create_count} created, {skipped_count} skipped."
  276. )
  277. try:
  278. click.echo("Creating dataset vector database index: {}".format(dataset.id))
  279. if dataset.index_struct_dict:
  280. if dataset.index_struct_dict["type"] == vector_type:
  281. skipped_count = skipped_count + 1
  282. continue
  283. collection_name = ""
  284. dataset_id = dataset.id
  285. if vector_type in upper_collection_vector_types:
  286. collection_name = Dataset.gen_collection_name_by_id(dataset_id)
  287. elif vector_type == VectorType.QDRANT:
  288. if dataset.collection_binding_id:
  289. dataset_collection_binding = (
  290. db.session.query(DatasetCollectionBinding)
  291. .filter(DatasetCollectionBinding.id == dataset.collection_binding_id)
  292. .one_or_none()
  293. )
  294. if dataset_collection_binding:
  295. collection_name = dataset_collection_binding.collection_name
  296. else:
  297. raise ValueError("Dataset Collection Binding not found")
  298. else:
  299. collection_name = Dataset.gen_collection_name_by_id(dataset_id)
  300. elif vector_type in lower_collection_vector_types:
  301. collection_name = Dataset.gen_collection_name_by_id(dataset_id).lower()
  302. else:
  303. raise ValueError(f"Vector store {vector_type} is not supported.")
  304. index_struct_dict = {"type": vector_type, "vector_store": {"class_prefix": collection_name}}
  305. dataset.index_struct = json.dumps(index_struct_dict)
  306. vector = Vector(dataset)
  307. click.echo(f"Migrating dataset {dataset.id}.")
  308. try:
  309. vector.delete()
  310. click.echo(
  311. click.style(f"Deleted vector index {collection_name} for dataset {dataset.id}.", fg="green")
  312. )
  313. except Exception as e:
  314. click.echo(
  315. click.style(
  316. f"Failed to delete vector index {collection_name} for dataset {dataset.id}.", fg="red"
  317. )
  318. )
  319. raise e
  320. dataset_documents = (
  321. db.session.query(DatasetDocument)
  322. .filter(
  323. DatasetDocument.dataset_id == dataset.id,
  324. DatasetDocument.indexing_status == "completed",
  325. DatasetDocument.enabled == True,
  326. DatasetDocument.archived == False,
  327. )
  328. .all()
  329. )
  330. documents = []
  331. segments_count = 0
  332. for dataset_document in dataset_documents:
  333. segments = (
  334. db.session.query(DocumentSegment)
  335. .filter(
  336. DocumentSegment.document_id == dataset_document.id,
  337. DocumentSegment.status == "completed",
  338. DocumentSegment.enabled == True,
  339. )
  340. .all()
  341. )
  342. for segment in segments:
  343. document = Document(
  344. page_content=segment.content,
  345. metadata={
  346. "doc_id": segment.index_node_id,
  347. "doc_hash": segment.index_node_hash,
  348. "document_id": segment.document_id,
  349. "dataset_id": segment.dataset_id,
  350. },
  351. )
  352. documents.append(document)
  353. segments_count = segments_count + 1
  354. if documents:
  355. try:
  356. click.echo(
  357. click.style(
  358. f"Creating vector index with {len(documents)} documents of {segments_count}"
  359. f" segments for dataset {dataset.id}.",
  360. fg="green",
  361. )
  362. )
  363. vector.create(documents)
  364. click.echo(click.style(f"Created vector index for dataset {dataset.id}.", fg="green"))
  365. except Exception as e:
  366. click.echo(click.style(f"Failed to created vector index for dataset {dataset.id}.", fg="red"))
  367. raise e
  368. db.session.add(dataset)
  369. db.session.commit()
  370. click.echo(f"Successfully migrated dataset {dataset.id}.")
  371. create_count += 1
  372. except Exception as e:
  373. db.session.rollback()
  374. click.echo(
  375. click.style("Error creating dataset index: {} {}".format(e.__class__.__name__, str(e)), fg="red")
  376. )
  377. continue
  378. click.echo(
  379. click.style(
  380. f"Migration complete. Created {create_count} dataset indexes. Skipped {skipped_count} datasets.", fg="green"
  381. )
  382. )
  383. @click.command("convert-to-agent-apps", help="Convert Agent Assistant to Agent App.")
  384. def convert_to_agent_apps():
  385. """
  386. Convert Agent Assistant to Agent App.
  387. """
  388. click.echo(click.style("Starting convert to agent apps.", fg="green"))
  389. proceeded_app_ids = []
  390. while True:
  391. # fetch first 1000 apps
  392. sql_query = """SELECT a.id AS id FROM apps a
  393. INNER JOIN app_model_configs am ON a.app_model_config_id=am.id
  394. WHERE a.mode = 'chat'
  395. AND am.agent_mode is not null
  396. AND (
  397. am.agent_mode like '%"strategy": "function_call"%'
  398. OR am.agent_mode like '%"strategy": "react"%'
  399. )
  400. AND (
  401. am.agent_mode like '{"enabled": true%'
  402. OR am.agent_mode like '{"max_iteration": %'
  403. ) ORDER BY a.created_at DESC LIMIT 1000
  404. """
  405. with db.engine.begin() as conn:
  406. rs = conn.execute(db.text(sql_query))
  407. apps = []
  408. for i in rs:
  409. app_id = str(i.id)
  410. if app_id not in proceeded_app_ids:
  411. proceeded_app_ids.append(app_id)
  412. app = db.session.query(App).filter(App.id == app_id).first()
  413. if app is not None:
  414. apps.append(app)
  415. if len(apps) == 0:
  416. break
  417. for app in apps:
  418. click.echo("Converting app: {}".format(app.id))
  419. try:
  420. app.mode = AppMode.AGENT_CHAT.value
  421. db.session.commit()
  422. # update conversation mode to agent
  423. db.session.query(Conversation).filter(Conversation.app_id == app.id).update(
  424. {Conversation.mode: AppMode.AGENT_CHAT.value}
  425. )
  426. db.session.commit()
  427. click.echo(click.style("Converted app: {}".format(app.id), fg="green"))
  428. except Exception as e:
  429. click.echo(click.style("Convert app error: {} {}".format(e.__class__.__name__, str(e)), fg="red"))
  430. click.echo(click.style("Conversion complete. Converted {} agent apps.".format(len(proceeded_app_ids)), fg="green"))
  431. @click.command("add-qdrant-index", help="Add Qdrant index.")
  432. @click.option("--field", default="metadata.doc_id", prompt=False, help="Index field , default is metadata.doc_id.")
  433. def add_qdrant_index(field: str):
  434. click.echo(click.style("Starting Qdrant index creation.", fg="green"))
  435. create_count = 0
  436. try:
  437. bindings = db.session.query(DatasetCollectionBinding).all()
  438. if not bindings:
  439. click.echo(click.style("No dataset collection bindings found.", fg="red"))
  440. return
  441. import qdrant_client
  442. from qdrant_client.http.exceptions import UnexpectedResponse
  443. from qdrant_client.http.models import PayloadSchemaType
  444. from core.rag.datasource.vdb.qdrant.qdrant_vector import QdrantConfig
  445. for binding in bindings:
  446. if dify_config.QDRANT_URL is None:
  447. raise ValueError("Qdrant URL is required.")
  448. qdrant_config = QdrantConfig(
  449. endpoint=dify_config.QDRANT_URL,
  450. api_key=dify_config.QDRANT_API_KEY,
  451. root_path=current_app.root_path,
  452. timeout=dify_config.QDRANT_CLIENT_TIMEOUT,
  453. grpc_port=dify_config.QDRANT_GRPC_PORT,
  454. prefer_grpc=dify_config.QDRANT_GRPC_ENABLED,
  455. )
  456. try:
  457. client = qdrant_client.QdrantClient(**qdrant_config.to_qdrant_params())
  458. # create payload index
  459. client.create_payload_index(binding.collection_name, field, field_schema=PayloadSchemaType.KEYWORD)
  460. create_count += 1
  461. except UnexpectedResponse as e:
  462. # Collection does not exist, so return
  463. if e.status_code == 404:
  464. click.echo(click.style(f"Collection not found: {binding.collection_name}.", fg="red"))
  465. continue
  466. # Some other error occurred, so re-raise the exception
  467. else:
  468. click.echo(
  469. click.style(
  470. f"Failed to create Qdrant index for collection: {binding.collection_name}.", fg="red"
  471. )
  472. )
  473. except Exception:
  474. click.echo(click.style("Failed to create Qdrant client.", fg="red"))
  475. click.echo(click.style(f"Index creation complete. Created {create_count} collection indexes.", fg="green"))
  476. @click.command("old-metadata-migration", help="Old metadata migration.")
  477. def old_metadata_migration():
  478. """
  479. Old metadata migration.
  480. """
  481. click.echo(click.style("Starting old metadata migration.", fg="green"))
  482. page = 1
  483. while True:
  484. try:
  485. documents = (
  486. DatasetDocument.query.filter(DatasetDocument.doc_metadata is not None)
  487. .order_by(DatasetDocument.created_at.desc())
  488. .paginate(page=page, per_page=50)
  489. )
  490. except NotFound:
  491. break
  492. if not documents:
  493. break
  494. for document in documents:
  495. if document.doc_metadata:
  496. doc_metadata = document.doc_metadata
  497. for key, value in doc_metadata.items():
  498. for field in BuiltInField:
  499. if field.value == key:
  500. break
  501. else:
  502. dataset_metadata = (
  503. db.session.query(DatasetMetadata)
  504. .filter(DatasetMetadata.dataset_id == document.dataset_id, DatasetMetadata.name == key)
  505. .first()
  506. )
  507. if not dataset_metadata:
  508. dataset_metadata = DatasetMetadata(
  509. tenant_id=document.tenant_id,
  510. dataset_id=document.dataset_id,
  511. name=key,
  512. type="string",
  513. created_by=document.created_by,
  514. )
  515. db.session.add(dataset_metadata)
  516. db.session.flush()
  517. dataset_metadata_binding = DatasetMetadataBinding(
  518. tenant_id=document.tenant_id,
  519. dataset_id=document.dataset_id,
  520. metadata_id=dataset_metadata.id,
  521. document_id=document.id,
  522. created_by=document.created_by,
  523. )
  524. db.session.add(dataset_metadata_binding)
  525. else:
  526. dataset_metadata_binding = DatasetMetadataBinding.query.filter(
  527. DatasetMetadataBinding.dataset_id == document.dataset_id,
  528. DatasetMetadataBinding.document_id == document.id,
  529. DatasetMetadataBinding.metadata_id == dataset_metadata.id,
  530. ).first()
  531. if not dataset_metadata_binding:
  532. dataset_metadata_binding = DatasetMetadataBinding(
  533. tenant_id=document.tenant_id,
  534. dataset_id=document.dataset_id,
  535. metadata_id=dataset_metadata.id,
  536. document_id=document.id,
  537. created_by=document.created_by,
  538. )
  539. db.session.add(dataset_metadata_binding)
  540. db.session.commit()
  541. page += 1
  542. click.echo(click.style("Old metadata migration completed.", fg="green"))
  543. @click.command("create-tenant", help="Create account and tenant.")
  544. @click.option("--email", prompt=True, help="Tenant account email.")
  545. @click.option("--name", prompt=True, help="Workspace name.")
  546. @click.option("--language", prompt=True, help="Account language, default: en-US.")
  547. def create_tenant(email: str, language: Optional[str] = None, name: Optional[str] = None):
  548. """
  549. Create tenant account
  550. """
  551. if not email:
  552. click.echo(click.style("Email is required.", fg="red"))
  553. return
  554. # Create account
  555. email = email.strip()
  556. if "@" not in email:
  557. click.echo(click.style("Invalid email address.", fg="red"))
  558. return
  559. account_name = email.split("@")[0]
  560. if language not in languages:
  561. language = "en-US"
  562. # Validates name encoding for non-Latin characters.
  563. name = name.strip().encode("utf-8").decode("utf-8") if name else None
  564. # generate random password
  565. new_password = secrets.token_urlsafe(16)
  566. # register account
  567. account = RegisterService.register(
  568. email=email,
  569. name=account_name,
  570. password=new_password,
  571. language=language,
  572. create_workspace_required=False,
  573. )
  574. TenantService.create_owner_tenant_if_not_exist(account, name)
  575. click.echo(
  576. click.style(
  577. "Account and tenant created.\nAccount: {}\nPassword: {}".format(email, new_password),
  578. fg="green",
  579. )
  580. )
  581. @click.command("upgrade-db", help="Upgrade the database")
  582. def upgrade_db():
  583. click.echo("Preparing database migration...")
  584. lock = redis_client.lock(name="db_upgrade_lock", timeout=60)
  585. if lock.acquire(blocking=False):
  586. try:
  587. click.echo(click.style("Starting database migration.", fg="green"))
  588. # run db migration
  589. import flask_migrate # type: ignore
  590. flask_migrate.upgrade()
  591. click.echo(click.style("Database migration successful!", fg="green"))
  592. except Exception:
  593. logging.exception("Failed to execute database migration")
  594. finally:
  595. lock.release()
  596. else:
  597. click.echo("Database migration skipped")
  598. @click.command("fix-app-site-missing", help="Fix app related site missing issue.")
  599. def fix_app_site_missing():
  600. """
  601. Fix app related site missing issue.
  602. """
  603. click.echo(click.style("Starting fix for missing app-related sites.", fg="green"))
  604. failed_app_ids = []
  605. while True:
  606. sql = """select apps.id as id from apps left join sites on sites.app_id=apps.id
  607. where sites.id is null limit 1000"""
  608. with db.engine.begin() as conn:
  609. rs = conn.execute(db.text(sql))
  610. processed_count = 0
  611. for i in rs:
  612. processed_count += 1
  613. app_id = str(i.id)
  614. if app_id in failed_app_ids:
  615. continue
  616. try:
  617. app = db.session.query(App).filter(App.id == app_id).first()
  618. if not app:
  619. print(f"App {app_id} not found")
  620. continue
  621. tenant = app.tenant
  622. if tenant:
  623. accounts = tenant.get_accounts()
  624. if not accounts:
  625. print("Fix failed for app {}".format(app.id))
  626. continue
  627. account = accounts[0]
  628. print("Fixing missing site for app {}".format(app.id))
  629. app_was_created.send(app, account=account)
  630. except Exception:
  631. failed_app_ids.append(app_id)
  632. click.echo(click.style("Failed to fix missing site for app {}".format(app_id), fg="red"))
  633. logging.exception(f"Failed to fix app related site missing issue, app_id: {app_id}")
  634. continue
  635. if not processed_count:
  636. break
  637. click.echo(click.style("Fix for missing app-related sites completed successfully!", fg="green"))
  638. @click.command("migrate-data-for-plugin", help="Migrate data for plugin.")
  639. def migrate_data_for_plugin():
  640. """
  641. Migrate data for plugin.
  642. """
  643. click.echo(click.style("Starting migrate data for plugin.", fg="white"))
  644. PluginDataMigration.migrate()
  645. click.echo(click.style("Migrate data for plugin completed.", fg="green"))
  646. @click.command("extract-plugins", help="Extract plugins.")
  647. @click.option("--output_file", prompt=True, help="The file to store the extracted plugins.", default="plugins.jsonl")
  648. @click.option("--workers", prompt=True, help="The number of workers to extract plugins.", default=10)
  649. def extract_plugins(output_file: str, workers: int):
  650. """
  651. Extract plugins.
  652. """
  653. click.echo(click.style("Starting extract plugins.", fg="white"))
  654. PluginMigration.extract_plugins(output_file, workers)
  655. click.echo(click.style("Extract plugins completed.", fg="green"))
  656. @click.command("extract-unique-identifiers", help="Extract unique identifiers.")
  657. @click.option(
  658. "--output_file",
  659. prompt=True,
  660. help="The file to store the extracted unique identifiers.",
  661. default="unique_identifiers.json",
  662. )
  663. @click.option(
  664. "--input_file", prompt=True, help="The file to store the extracted unique identifiers.", default="plugins.jsonl"
  665. )
  666. def extract_unique_plugins(output_file: str, input_file: str):
  667. """
  668. Extract unique plugins.
  669. """
  670. click.echo(click.style("Starting extract unique plugins.", fg="white"))
  671. PluginMigration.extract_unique_plugins_to_file(input_file, output_file)
  672. click.echo(click.style("Extract unique plugins completed.", fg="green"))
  673. @click.command("install-plugins", help="Install plugins.")
  674. @click.option(
  675. "--input_file", prompt=True, help="The file to store the extracted unique identifiers.", default="plugins.jsonl"
  676. )
  677. @click.option(
  678. "--output_file", prompt=True, help="The file to store the installed plugins.", default="installed_plugins.jsonl"
  679. )
  680. @click.option("--workers", prompt=True, help="The number of workers to install plugins.", default=100)
  681. def install_plugins(input_file: str, output_file: str, workers: int):
  682. """
  683. Install plugins.
  684. """
  685. click.echo(click.style("Starting install plugins.", fg="white"))
  686. PluginMigration.install_plugins(input_file, output_file, workers)
  687. click.echo(click.style("Install plugins completed.", fg="green"))