Add File

backend/apps/system/crud/workspace.py

@@ -0,0 +1,69 @@
+
+from collections import defaultdict
+from typing import Optional
+from sqlmodel import Session, func, select, update
+
+from apps.system.models.system_model import UserWsModel
+from apps.system.models.user import UserModel
+
+async def reset_single_user_oid(session: Session, uid: int, oid: int, add: Optional[bool] = True):
+    user_model = session.get(UserModel, uid)
+    if not user_model:
+        return
+    origin_oid = user_model.oid
+    if add and (not origin_oid or origin_oid == 0):
+        user_model.oid = oid
+        session.add(user_model)
+    if not add and origin_oid and origin_oid == oid:
+        user_model.oid = 0
+        user_ws = session.exec(select(UserWsModel).where(UserWsModel.uid == uid, UserWsModel.oid != oid)).first()
+        if user_ws:
+            user_model.oid = user_ws.oid
+        session.add(user_model)
+    
+async def reset_user_oid(session: Session, oid: int):
+    stmt = (
+        select(
+            UserModel.id,
+            UserModel.oid,
+            UserWsModel.oid.label("associated_oid")
+        )
+        .join(UserWsModel, UserModel.id == UserWsModel.uid, isouter=True)
+        .where(UserModel.id != 1)
+    )
+    
+    user_filter = (
+        select(UserModel.id)
+            .join(UserWsModel, UserModel.id == UserWsModel.uid)
+            .where(UserWsModel.oid == oid)
+            .distinct()
+    )
+    stmt = stmt.where(UserModel.id.in_(user_filter))
+    
+    result_user_list = session.exec(stmt).all()
+    if not result_user_list:
+        return
+    
+    merged = defaultdict(list)
+    extra_attrs = {}
+
+    for (id, oid, associated_oid) in result_user_list:
+        item = {"id": id, "oid": oid}
+        merged[id].append(associated_oid)
+        if id not in extra_attrs:
+            extra_attrs[id] = {k: v for k, v in item.items()}
+
+    # 组合结果
+    result = [
+        {**extra_attrs[user_id], "oid_list": oid_list} 
+        for user_id, oid_list in merged.items()
+    ]
+    
+    for row in result:
+        origin_oid = row['oid']
+        oid_list: list = list(filter(lambda x: x != oid, row['oid_list']))
+        if origin_oid not in oid_list:
+            row['oid'] = oid_list[0] if oid_list else 0
+            if row['oid'] != origin_oid:
+                update_stmt = update(UserModel).where(UserModel.id == row['id']).values(oid=row['oid'])
+                session.exec(update_stmt)