self-constructing_graph/agent/libs/user_data_relation.py

from datetime import datetime
from sqlalchemy.orm import Session
from ..models.db.graph import DbUserDataRelation

class UserDataRelationBusiness:
    def __init__(self, db: Session):
        self.db = db
    
    def create_relation(self, user_id: int, data_category: str, data_id: int, role_id: int):
        """
        创建用户数据关联
        :param user_id: 用户ID
        :param data_category: 数据类别(表名)
        :param data_id: 数据ID
        :param role_id: 角色ID
        :return: DbUserDataRelation对象
        """
        relation = DbUserDataRelation(
            user_id=user_id,
            data_category=data_category,
            data_id=data_id,
            role_id=role_id,
            created=datetime.now(),
            updated=datetime.now()
        )
        self.db.add(relation)
        self.db.commit()
        self.db.refresh(relation)
        return relation
    
    def get_relation(self, relation_id: int):
        """
        获取关联关系
        :param relation_id: 关联ID
        :return: DbUserDataRelation对象
        """
        return self.db.query(DbUserDataRelation).filter(DbUserDataRelation.id == relation_id).first()
    
    def get_user_relations(self, user_id: int):
        """
        获取用户的所有关联关系
        :param user_id: 用户ID
        :return: DbUserDataRelation对象列表
        """
        return self.db.query(DbUserDataRelation).filter(DbUserDataRelation.user_id == user_id).all()
    
    def get_data_relations(self, data_category: str, data_id: int):
        """
        获取特定数据的所有关联关系
        :param data_category: 数据类别(表名)
        :param data_id: 数据ID
        :return: DbUserDataRelation对象列表
        """
        return self.db.query(DbUserDataRelation).filter(
            DbUserDataRelation.data_category == data_category,
            DbUserDataRelation.data_id == data_id
        ).all()
    
    def update_relation(self, relation_id: int, role_id: int):
        """
        更新关联关系的角色ID
        :param relation_id: 关联ID
        :param role_id: 新的角色ID
        :return: 更新后的DbUserDataRelation对象
        """
        relation = self.get_relation(relation_id)
        if relation:
            relation.role_id = role_id
            relation.updated = datetime.now()
            self.db.commit()
            self.db.refresh(relation)
        return relation
    
    def delete_relation(self, relation_id: int):
        """
        删除关联关系
        :param relation_id: 关联ID
        :return: 是否删除成功
        """
        relation = self.get_relation(relation_id)
        if relation:
            self.db.delete(relation)
            self.db.commit()
            return True
        return False
    
    def delete_user_relations(self, user_id: int):
        """
        删除用户的所有关联关系
        :param user_id: 用户ID
        :return: 删除的记录数
        """
        count = self.db.query(DbUserDataRelation).filter(DbUserDataRelation.user_id == user_id).delete()
        self.db.commit()
        return count
    
    def delete_data_relations(self, data_category: str, data_id: int):
        """
        删除特定数据的所有关联关系
        :param data_category: 数据类别(表名)
        :param data_id: 数据ID
        :return: 删除的记录数
        """
        count = self.db.query(DbUserDataRelation).filter(
            DbUserDataRelation.data_category == data_category,
            DbUserDataRelation.data_id == data_id
        ).delete()
        self.db.commit()
        return count