82 lines
2.6 KiB
Python
82 lines
2.6 KiB
Python
import hashlib
|
|
from sqlalchemy import Column, String
|
|
from sqlalchemy.orm import declarative_base, declared_attr
|
|
import ulid
|
|
|
|
Base = declarative_base()
|
|
|
|
# Base62 字符集
|
|
ALPHABET = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
|
|
|
|
|
|
def encode_base62(num: int, length: int = None) -> str:
|
|
"""将一个整数编码为 Base62 字符串，支持左侧补 0 固定长度"""
|
|
if num == 0:
|
|
res = ALPHABET[0]
|
|
else:
|
|
arr = []
|
|
base = len(ALPHABET)
|
|
while num:
|
|
num, rem = divmod(num, base)
|
|
arr.append(ALPHABET[rem])
|
|
res = "".join(reversed(arr))
|
|
|
|
if length:
|
|
# 如果太长就截取，太短就用第一个字符(0)补齐
|
|
return res.rjust(length, ALPHABET[0])[-length:]
|
|
return res
|
|
|
|
|
|
def encode_tablename_to_prefix(tablename: str, length: int) -> str:
|
|
"""
|
|
将表名 encode 成固定长度的前缀，使用 MD5 拿到表名的唯一哈希值，转为整数后用 Base62 截取前 N 位
|
|
"""
|
|
# 计算表名的 MD5 哈希
|
|
hasher = hashlib.md5(tablename.encode("utf-8"))
|
|
hex_digest = hasher.hexdigest()
|
|
|
|
# 将十六进制哈希转为大整数
|
|
num = int(hex_digest, 16)
|
|
|
|
# 转为 Base62 字符串
|
|
return encode_base62(num, length=length)
|
|
|
|
|
|
def gen_convex_payload(length: int) -> str:
|
|
"""生成固定长度的 ULID Payload"""
|
|
u = ulid.new()
|
|
return encode_base62(u.int, length=length)
|
|
|
|
|
|
# --- 注入器类 ---
|
|
|
|
class AutoID:
|
|
"""
|
|
全自动 Convex 风格 ID 注入器
|
|
不再需要手动定义任何 prefix，全自动根据 __tablename__ 编码生成
|
|
支持子类自定义长度
|
|
"""
|
|
# 默认配置（如果子类不定义，则使用这些默认值）
|
|
_id_prefix_len = 3 # 前缀长度
|
|
_id_payload_len = 22 # 主体长度（ULID 转 Base62 最大约 22 位）
|
|
|
|
@classmethod
|
|
def id_length(cls) -> int:
|
|
"""返回当前类 ID 的总长度"""
|
|
return cls._id_prefix_len + cls._id_payload_len
|
|
|
|
@declared_attr
|
|
def id(cls):
|
|
# 运行时动态获取子类的 __tablename__
|
|
tablename = cls.__tablename__
|
|
# 获取子类配置的长度
|
|
p_len = cls._id_prefix_len
|
|
s_len = cls._id_payload_len
|
|
prefix = encode_tablename_to_prefix(tablename, length=p_len)
|
|
|
|
# 核心工厂函数：拼接【编码后的表前缀】和【递增Payload】
|
|
def id_factory():
|
|
payload = gen_convex_payload(length=s_len)
|
|
return f"{prefix}{payload}"
|
|
|
|
return Column(String(cls.id_length()), primary_key=True, default=id_factory) |