RagFlow ORM

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

api/db/
├── __init__.py          # 枚举定义和常量
├── db_models.py         # 数据库模型定义
├── db_utils.py          # 数据库工具函数
├── runtime_config.py    # 运行时配置管理
├── init_data.py         # 初始化数据
└── services/            # 业务服务层
    ├── user_service.py
    ├── document_service.py
    ├── knowledgebase_service.py
    ├── canvas_service.py
    └── ... (其他服务)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

from peewee import Model, CharField, IntegerField
from peewee import SqliteDatabase

db = SqliteDatabase('my_database.db')

# 定义模型（表）
class User(Model):
    username = CharField(max_length=50)
    age = IntegerField()

    class Meta:
        database = db  # 指定数据库
        indexes = ((('username',), True),)  # 唯一索引

# 事务
with db.atomic():
    User.create(username='Bob', age=25)

1
2
3
4

from playhouse.shortcuts import model_to_dict

user = User.get(User.id == 1)
data = model_to_dict(user)

1
2
3
4
5
6
7

from playhouse.migrate import SqliteMigrator, migrate
from peewee import CharField

migrator = SqliteMigrator(db)
migrate(
    migrator.add_column('user', 'email', CharField(default=''))
)

1
2
3

from playhouse.pool import PooledPostgresqlDatabase

db = PooledPostgresqlDatabase('mydb', max_connections=8, user='user', password='pass')

1
2
3
4
5
6
7
8

from playhouse.signals import Model, post_save

class User(Model):
    username = CharField()

@post_save(sender=User)
def after_save_handler(model_class, instance, created):
    print(f"{instance.username} saved, created={created}")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

from peewee import Field
import json

class JSONField(Field):
    field_type = 'TEXT'  # 数据库存储类型

    def db_value(self, value):
        # Python 对象 → 数据库字符串
        if value is not None:
            return json.dumps(value)
        return None

    def python_value(self, value):
        # 数据库字符串 → Python 对象
        if value is not None:
            return json.loads(value)
        return None

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

from peewee import Model, SqliteDatabase

db = SqliteDatabase('test.db')

class Document(Model):
    data = JSONField()

    class Meta:
        database = db

db.connect()
db.create_tables([Document])

# 测试
doc = Document.create(data={"name": "Alice", "age": 30})
retrieved = Document.get(Document.id == doc.id)
print(retrieved.data)  # {'name': 'Alice', 'age': 30}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

from peewee import Field
from cryptography.fernet import Fernet

key = Fernet.generate_key()
cipher = Fernet(key)

class EncryptedField(Field):
    field_type = 'TEXT'

    def db_value(self, value):
        if value is not None:
            return cipher.encrypt(value.encode()).decode()
        return None

    def python_value(self, value):
        if value is not None:
            return cipher.decrypt(value.encode()).decode()
        return None

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71

# 1. 根据数据，选择长文本类型
class TextFieldType(Enum):
    MYSQL = "LONGTEXT"
    POSTGRES = "TEXT"


class LongTextField(TextField):
    field_type = TextFieldType[settings.DATABASE_TYPE.upper()].value


# 2. 自定义 JSON 和 List 类型
class JSONField(LongTextField):
    default_value = {}

    def __init__(self, object_hook=None, object_pairs_hook=None, **kwargs):
        self._object_hook = object_hook
        self._object_pairs_hook = object_pairs_hook
        super().__init__(**kwargs)

    def db_value(self, value):
        if value is None:
            value = self.default_value
        return utils.json_dumps(value)

    def python_value(self, value):
        if not value:
            return self.default_value
        return utils.json_loads(value, object_hook=self._object_hook, object_pairs_hook=self._object_pairs_hook)


class ListField(JSONField):
    default_value = []


# 3. 自定义序列化字段，可选序列化方式，支持 json 和 pickle
# 相比于 JSONField，添加了 object_hook
class SerializedType(IntEnum):
    PICKLE = 1
    JSON = 2

class SerializedField(LongTextField):
    def __init__(self, serialized_type=SerializedType.PICKLE, object_hook=None, object_pairs_hook=None, **kwargs):
        self._serialized_type = serialized_type
        self._object_hook = object_hook
        self._object_pairs_hook = object_pairs_hook
        super().__init__(**kwargs)

    def db_value(self, value):
        if self._serialized_type == SerializedType.PICKLE:
            return utils.serialize_b64(value, to_str=True)
        elif self._serialized_type == SerializedType.JSON:
            if value is None:
                return None
            return utils.json_dumps(value, with_type=True)
        else:
            raise ValueError(f"the serialized type {self._serialized_type} is not supported")

    def python_value(self, value):
        if self._serialized_type == SerializedType.PICKLE:
            return utils.deserialize_b64(value)
        elif self._serialized_type == SerializedType.JSON:
            if value is None:
                return {}
            return utils.json_loads(value, object_hook=self._object_hook, object_pairs_hook=self._object_pairs_hook)
        else:
            raise ValueError(f"the serialized type {self._serialized_type} is not supported")


class JsonSerializedField(SerializedField):
    def __init__(self, object_hook=utils.from_dict_hook, object_pairs_hook=None, **kwargs):
        super(JsonSerializedField, self).__init__(serialized_type=SerializedType.JSON, object_hook=object_hook, object_pairs_hook=object_pairs_hook, **kwargs)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

import json

class User:
    def __init__(self, name, age):
        self.name = name
        self.age = age
    def __repr__(self):
        return f"User(name={self.name}, age={self.age})"

def user_hook(d):
    # 把 dict 转换成 User 对象
    if "name" in d and "age" in d:
        return User(d["name"], d["age"])
    return d

data = '{"name": "Alice", "age": 30}'
user = json.loads(data, object_hook=user_hook)
print(user)  # User(name=Alice, age=30)

1
2
3
4
5
6
7
8

import json
from collections import OrderedDict

data = '{"b": 2, "a": 1, "c": 3}'

# 使用 object_pairs_hook 保留顺序
ordered = json.loads(data, object_pairs_hook=OrderedDict)
print(ordered)  # OrderedDict([('b', 2), ('a', 1), ('c', 3)])

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128

class BaseModel(Model):
    # 定义基础字段：创建时间、创建日期、更新时间、更新日期，并建立索引
    create_time = BigIntegerField(null=True, index=True)  # 时间戳形式的创建时间
    create_date = DateTimeField(null=True, index=True)    # datetime 对象形式的创建时间
    update_time = BigIntegerField(null=True, index=True)  # 时间戳形式的更新时间
    update_date = DateTimeField(null=True, index=True)    # datetime 对象形式的更新时间

    # 将模型转为 JSON（弃用，内部调用 to_dict）
    def to_json(self):
        return self.to_dict()

    # 将模型数据转换为字典（直接返回 __data__，包含字段名和对应值）
    def to_dict(self):
        return self.__dict__["__data__"]

    # 返回人类可读的模型字典，可选择只包含主键和指定字段
    def to_human_model_dict(self, only_primary_with: list = None):
        model_dict = self.__dict__["__data__"]

        if not only_primary_with:
            # 移除字段前缀，返回所有字段
            return {remove_field_name_prefix(k): v for k, v in model_dict.items()}

        human_model_dict = {}
        # 添加主键字段
        for k in self._meta.primary_key.field_names:
            human_model_dict[remove_field_name_prefix(k)] = model_dict[k]
        # 添加额外指定字段
        for k in only_primary_with:
            human_model_dict[k] = model_dict[f"f_{k}"]
        return human_model_dict

    # 返回模型的元数据对象（_meta）
    @property
    def meta(self) -> Metadata:
        return self._meta

    # 获取主键字段名列表
    @classmethod
    def get_primary_keys_name(cls):
        if isinstance(cls._meta.primary_key, CompositeKey):
            return cls._meta.primary_key.field_names
        else:
            return [cls._meta.primary_key.name]

    # 根据属性名获取字段对象（用于动态访问字段）
    @classmethod
    def getter_by(cls, attr):
        return operator.attrgetter(attr)(cls)

    # 构建查询方法，可按字段值过滤，并支持范围、列表、多条件查询
    @classmethod
    def query(cls, reverse=None, order_by=None, **kwargs):
        filters = []
        for f_n, f_v in kwargs.items():
            attr_name = "%s" % f_n
            # 忽略不存在的字段或值为 None 的条件
            if not hasattr(cls, attr_name) or f_v is None:
                continue
            # 处理列表/集合类型的查询值
            if type(f_v) in {list, set}:
                f_v = list(f_v)
                # 对连续类型字段（如时间戳、数值）进行范围查询
                if is_continuous_field(type(getattr(cls, attr_name))):
                    if len(f_v) == 2:
                        # 如果是日期字符串，将其转换为时间戳
                        for i, v in enumerate(f_v):
                            if isinstance(v, str) and f_n in auto_date_timestamp_field():
                                f_v[i] = utils.date_string_to_timestamp(v)
                        lt_value, gt_value = f_v
                        # 构建 between 或 >= / <= 查询条件
                        if lt_value is not None and gt_value is not None:
                            filters.append(cls.getter_by(attr_name).between(lt_value, gt_value))
                        elif lt_value is not None:
                            filters.append(operator.attrgetter(attr_name)(cls) >= lt_value)
                        elif gt_value is not None:
                            filters.append(operator.attrgetter(attr_name)(cls) <= gt_value)
                else:
                    # 非连续类型字段使用 IN 查询
                    filters.append(operator.attrgetter(attr_name)(cls) << f_v)
            else:
                # 单值字段使用等于查询
                filters.append(operator.attrgetter(attr_name)(cls) == f_v)
        # 构建查询结果
        if filters:
            query_records = cls.select().where(*filters)
            # 支持排序
            if reverse is not None:
                if not order_by or not hasattr(cls, f"{order_by}"):
                    order_by = "create_time"
                if reverse is True:
                    query_records = query_records.order_by(cls.getter_by(f"{order_by}").desc())
                elif reverse is False:
                    query_records = query_records.order_by(cls.getter_by(f"{order_by}").asc())
            return [query_record for query_record in query_records]
        else:
            return []

    # 插入记录时自动填充 create_time
    @classmethod
    def insert(cls, __data=None, **insert):
        if isinstance(__data, dict) and __data:
          # __data 保存的是 Feild -> 值的映射
            __data[cls._meta.combined["create_time"]] = utils.current_timestamp()
        if insert:
            # insert 是 str key -> 值的映射
            insert["create_time"] = utils.current_timestamp()
        return super().insert(__data, **insert)

    # 数据标准化方法，供 insert/update 调用，自动更新时间和日期字段
    @classmethod
    def _normalize_data(cls, data, kwargs):
        normalized = super()._normalize_data(data, kwargs)
        if not normalized:
            return {}
        # 自动更新 update_time
        normalized[cls._meta.combined["update_time"]] = utils.current_timestamp()

        # 自动同步 *_time 和 *_date 字段
        for f_n in AUTO_DATE_TIMESTAMP_FIELD_PREFIX:
            # 判断字段是否存在，并且时间字段不为空
            if {f"{f_n}_time", f"{f_n}_date"}.issubset(cls._meta.combined.keys()) and \
               cls._meta.combined[f"{f_n}_time"] in normalized and \
               normalized[cls._meta.combined[f"{f_n}_time"]] is not None:
                normalized[cls._meta.combined[f"{f_n}_date"]] = utils.timestamp_to_date(
                    normalized[cls._meta.combined[f"{f_n}_time"]]
                )
        return normalized

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35

# 1. 定义数据库连接池可选枚举
class PooledDatabase(Enum):
    MYSQL = RetryingPooledMySQLDatabase
    POSTGRES = PooledPostgresqlDatabase


# 2. singleton 保证进程内数据库唯一
@singleton
class BaseDataBase:
    def __init__(self):
        # 通过配置初始化数据库连接池
        database_config = settings.DATABASE.copy()
        db_name = database_config.pop("name")
        self.database_connection = PooledDatabase[settings.DATABASE_TYPE.upper()].value(db_name, **database_config)
        logging.info("init database on cluster mode successfully")

# 3. 实例化全局数据库链接对象
DB = BaseDataBase().database_connection

# 4. 所有 model 继承自 DataBaseModel，持有数据库链接对象
class DataBaseModel(BaseModel):
    class Meta:
        database = DB

def singleton(cls, *args, **kw):
    instances = {}

    def _singleton():
        # 单例模式，通过类名和进程 ID 作为键，保证进程内唯一
        key = str(cls) + str(os.getpid())
        if key not in instances:
            instances[key] = cls(*args, **kw)
        return instances[key]

    return _singleton

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43

class RetryingPooledMySQLDatabase(PooledMySQLDatabase):
    def __init__(self, *args, **kwargs):
        self.max_retries = kwargs.pop('max_retries', 5)
        self.retry_delay = kwargs.pop('retry_delay', 1)
        super().__init__(*args, **kwargs)

    def execute_sql(self, sql, params=None, commit=True):
        from peewee import OperationalError
        for attempt in range(self.max_retries + 1):
            try:
                return super().execute_sql(sql, params, commit)
            except OperationalError as e:
                if e.args[0] in (2013, 2006) and attempt < self.max_retries:
                    logging.warning(
                        f"Lost connection (attempt {attempt+1}/{self.max_retries}): {e}"
                    )
                    self._handle_connection_loss()
                    time.sleep(self.retry_delay * (2 ** attempt))
                else:
                    logging.error(f"DB execution failure: {e}")
                    raise
        return None

    def _handle_connection_loss(self):
        # 连接丢失时，关闭所有连接并重新连接
        # 这个不会对其他正在使用中的链接有影响么？
        self.close_all()
        self.connect()

    def begin(self):
        from peewee import OperationalError
        for attempt in range(self.max_retries + 1):
            try:
                return super().begin()
            except OperationalError as e:
                if e.args[0] in (2013, 2006) and attempt < self.max_retries:
                    logging.warning(
                        f"Lost connection during transaction (attempt {attempt+1}/{self.max_retries})"
                    )
                    self._handle_connection_loss()
                    time.sleep(self.retry_delay * (2 ** attempt))
                else:
                    raise

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100

class DatabaseLock(Enum):
    MYSQL = MysqlDatabaseLock
    POSTGRES = PostgresDatabaseLock

DB = BaseDataBase().database_connection
DB.lock = DatabaseLock[settings.DATABASE_TYPE.upper()].value


class MysqlDatabaseLock:
    def __init__(self, lock_name, timeout=10, db=None):
        self.lock_name = lock_name
        self.timeout = int(timeout)
        self.db = db if db else DB

    @with_retry(max_retries=3, retry_delay=1.0)
    def lock(self):
        # SQL parameters only support %s format placeholders
        cursor = self.db.execute_sql("SELECT GET_LOCK(%s, %s)", (self.lock_name, self.timeout))
        ret = cursor.fetchone()
        if ret[0] == 0:
            raise Exception(f"acquire mysql lock {self.lock_name} timeout")
        elif ret[0] == 1:
            return True
        else:
            raise Exception(f"failed to acquire lock {self.lock_name}")

    @with_retry(max_retries=3, retry_delay=1.0)
    def unlock(self):
        cursor = self.db.execute_sql("SELECT RELEASE_LOCK(%s)", (self.lock_name,))
        ret = cursor.fetchone()
        if ret[0] == 0:
            raise Exception(f"mysql lock {self.lock_name} was not established by this thread")
        elif ret[0] == 1:
            return True
        else:
            raise Exception(f"mysql lock {self.lock_name} does not exist")

    def __enter__(self):
        if isinstance(self.db, PooledMySQLDatabase):
            self.lock()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if isinstance(self.db, PooledMySQLDatabase):
            self.unlock()

    def __call__(self, func):
        @wraps(func)
        def magic(*args, **kwargs):
            with self:
                return func(*args, **kwargs)

        return magic


class PostgresDatabaseLock:
    def __init__(self, lock_name, timeout=10, db=None):
        self.lock_name = lock_name
        self.lock_id = int(hashlib.md5(lock_name.encode()).hexdigest(), 16) % (2**31 - 1)
        self.timeout = int(timeout)
        self.db = db if db else DB

    @with_retry(max_retries=3, retry_delay=1.0)
    def lock(self):
        cursor = self.db.execute_sql("SELECT pg_try_advisory_lock(%s)", (self.lock_id,))
        ret = cursor.fetchone()
        if ret[0] == 0:
            raise Exception(f"acquire postgres lock {self.lock_name} timeout")
        elif ret[0] == 1:
            return True
        else:
            raise Exception(f"failed to acquire lock {self.lock_name}")

    @with_retry(max_retries=3, retry_delay=1.0)
    def unlock(self):
        cursor = self.db.execute_sql("SELECT pg_advisory_unlock(%s)", (self.lock_id,))
        ret = cursor.fetchone()
        if ret[0] == 0:
            raise Exception(f"postgres lock {self.lock_name} was not established by this thread")
        elif ret[0] == 1:
            return True
        else:
            raise Exception(f"postgres lock {self.lock_name} does not exist")

    def __enter__(self):
        if isinstance(self.db, PooledPostgresqlDatabase):
            self.lock()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if isinstance(self.db, PooledPostgresqlDatabase):
            self.unlock()

    def __call__(self, func):
        @wraps(func)
        def magic(*args, **kwargs):
            with self:
                return func(*args, **kwargs)

        return magic

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38

# 给函数增加数据库上下文管理和分布式锁
@DB.connection_context()            # 确保函数运行时自动获取和释放数据库连接
@DB.lock("init_database_tables", 60)  # 获取名为 "init_database_tables" 的数据库锁，超时时间 60 秒
def init_database_tables(alter_fields=[]):
    # 获取当前模块中所有类（class）的成员
    members = inspect.getmembers(sys.modules[__name__], inspect.isclass)

    table_objs = []         # 存放需要处理的数据库表模型类
    create_failed_list = [] # 记录创建失败的表名

    # 遍历模块中定义的所有类
    for name, obj in members:
        # 只处理继承自 DataBaseModel 的子类（即用户自定义的 ORM 模型类）
        if obj != DataBaseModel and issubclass(obj, DataBaseModel):
            table_objs.append(obj)

            # 判断表是否已经存在
            if not obj.table_exists():
                logging.debug(f"start create table {obj.__name__}")
                try:
                    # 创建表，如果表已存在则不会抛错（safe=True）
                    obj.create_table(safe=True)
                    logging.debug(f"create table success: {obj.__name__}")
                except Exception as e:
                    # 如果建表失败，记录错误并加入失败列表
                    logging.exception(e)
                    create_failed_list.append(obj.__name__)
            else:
                # 如果表已经存在则跳过
                logging.debug(f"table {obj.__name__} already exists, skip creation.")

    # 如果有表创建失败，记录日志并抛出异常
    if create_failed_list:
        logging.error(f"create tables failed: {create_failed_list}")
        raise Exception(f"create tables failed: {create_failed_list}")

    # 调用数据库迁移函数（用于更新已有表的字段等结构变更）
    migrate_db()