From fe30e85ad48a5fd787940cf9d58fb87a12b0b562 Mon Sep 17 00:00:00 2001
From: zhaoqingang <zhaoqg0118@163.com>
Date: 星期四, 26 十二月 2024 10:28:55 +0800
Subject: [PATCH] 注册用户权限同步

---
 app/task/fetch_agent.py |   37 ++++++++++++++++++++++++++++++++-----
 1 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/app/task/fetch_agent.py b/app/task/fetch_agent.py
index 303d8f6..8507bee 100644
--- a/app/task/fetch_agent.py
+++ b/app/task/fetch_agent.py
@@ -62,6 +62,15 @@
     status = Column(String(1))  # 鐘舵��
     doc_num = Column(Integer)  # 鏂囨。
 
+
+class RgUserTenant(Base):
+    __tablename__ = 'user_tenant'
+    id = Column(String(36), primary_key=True)  # id
+    tenant_id = Column(String(32))  # 鍚嶇О
+    user_id = Column(String(32))
+    role = Column(String(32))  # 鍒涘缓浜篿d
+
+
 # 瑙f瀽鍚嶅瓧
 def parse_names(names_str: str) -> List[str]:
     return [name.strip() for name in names_str.split(',')]
@@ -228,9 +237,9 @@
     try:
         if names:
             query = db.query(Dialog.id, Dialog.name, Dialog.description, Dialog.status, Dialog.tenant_id) \
-                .filter( Dialog.name.in_(names))
+                .filter( Dialog.name.in_(names), Dialog.status == "1")
         else:
-            query = db.query(Dialog.id, Dialog.name, Dialog.description, Dialog.status, Dialog.tenant_id)
+            query = db.query(Dialog.id, Dialog.name, Dialog.description, Dialog.status, Dialog.tenant_id).filter( Dialog.status == "1")
 
         results = query.all()
         formatted_results = [
@@ -271,7 +280,7 @@
             if existing_agent:
                 existing_agent.name = row["name"]
                 existing_agent.description = row["description"]
-                existing_agent.status = row["status"]
+                # existing_agent.status = row["status"]
                 existing_agent.mode = row["mode"]
                 # existing_agent.tenant_id = get_rag_user_id(db, row["user_id"], type_dict[dialog_type])
             else:
@@ -324,6 +333,8 @@
         print("v2 Agents synchronized successfully")
     except Exception as e:
         print(f"v2 Failed to sync agents: {str(e)}")
+    finally:
+        db.close()
 
 def update_ids_in_local_knowledge(data, klg_type):
     type_dict = {"1": RAGFLOW, "2": BISHENG, "4": DIFY}
@@ -386,8 +397,24 @@
         print("sync knowledge successfully")
     except Exception as e:
         print(f"Failed to sync knowledge: {str(e)}")
+    finally:
+        db.close()
+
+
+def update_ragflow_user_tenant(user_id: str):
+    db = SessionRagflow()
+    try:
+        if user_id:
+            db.query(RgUserTenant).filter( RgUserTenant.user_id==user_id, RgUserTenant.role=="invite").update({"role": "normal"})
+            db.query(RgUserTenant).filter( RgUserTenant.tenant_id==user_id, RgUserTenant.role=="invite").update({"role": "normal"})
+        else:
+            db.query(RgUserTenant).filter(RgUserTenant.role=="invite").update({"role": "normal"})
+        db.commit()
+    finally:
+        db.close()
 
 
 if __name__ == "__main__":
-    a = get_data_from_dify_v2([])
-    print(a)
+    # a = get_data_from_dify_v2([])
+    # print(a)
+    update_ragflow_user_tenant("")

--
Gitblit v1.8.0