133 lines
6.3 KiB
Python
133 lines
6.3 KiB
Python
# requires: pymongo
|
||
"""MongoDB：AI 客服会话存储——Schema 校验、复合索引与 explain、聚合统计、幂等 upsert、事务、变更流。
|
||
|
||
运行：MONGO_URL="mongodb://127.0.0.1:27017/?directConnection=true" python mongo_demo.py
|
||
需要 MongoDB 以副本集方式运行（事务和变更流都要求副本集；单节点副本集即可）。没有设置环境变量时跳过。
|
||
使用一个随机名字的临时数据库，结束后删除。
|
||
"""
|
||
import os
|
||
import sys
|
||
import uuid
|
||
from datetime import datetime, timedelta, timezone
|
||
|
||
from pymongo import ASCENDING, DESCENDING, MongoClient
|
||
from pymongo.errors import DuplicateKeyError, WriteError
|
||
|
||
url = os.environ.get("MONGO_URL")
|
||
if not url:
|
||
print("未设置 MONGO_URL，跳过")
|
||
sys.exit(0)
|
||
|
||
client = MongoClient(url, serverSelectionTimeoutMS=3000)
|
||
db = client[f"learn_{uuid.uuid4().hex[:8]}"]
|
||
|
||
|
||
def plan_stages(plan: dict) -> set[str]:
|
||
"""递归收集 explain 结果中的执行阶段名（IXSCAN / COLLSCAN / FETCH ...）。"""
|
||
out = {plan.get("stage", "")}
|
||
for key in ("inputStage", "queryPlan"):
|
||
if key in plan:
|
||
out |= plan_stages(plan[key])
|
||
for child in plan.get("inputStages", []):
|
||
out |= plan_stages(child)
|
||
return out
|
||
|
||
|
||
try:
|
||
# 1. Schema 校验：文档模型灵活，但关键字段仍要约束
|
||
db.create_collection("conversations", validator={"$jsonSchema": {
|
||
"bsonType": "object",
|
||
"required": ["tenant_id", "user_id", "status", "created_at", "messages"],
|
||
"properties": {
|
||
"tenant_id": {"bsonType": "string"},
|
||
"status": {"enum": ["open", "handoff", "closed"]},
|
||
"messages": {"bsonType": "array", "maxItems": 200}, # 嵌入数组要设上限，防止文档无限增长（单文档上限 16MB）
|
||
},
|
||
}})
|
||
conv = db.conversations
|
||
try:
|
||
conv.insert_one({"tenant_id": "t1", "user_id": "u1", "status": "unknown", "created_at": datetime.now(timezone.utc), "messages": []})
|
||
raise AssertionError("非法 status 应该被拒绝")
|
||
except WriteError:
|
||
print("1. Schema 校验：非法 status 被拒绝")
|
||
|
||
# 2. 写入数据：消息嵌入在会话文档中（一起读、一起写，数量有上限）
|
||
now = datetime.now(timezone.utc)
|
||
docs = []
|
||
for i in range(300):
|
||
tenant = f"t{i % 3}"
|
||
docs.append({
|
||
"tenant_id": tenant, "user_id": f"u{i % 50}",
|
||
"status": ["open", "handoff", "closed"][i % 3],
|
||
"created_at": now - timedelta(minutes=i),
|
||
"channel": "web" if i % 2 else "app",
|
||
"messages": [{"role": "user", "text": f"问题{i}"}, {"role": "assistant", "text": f"回答{i}", "tokens": 20 + i % 7}],
|
||
})
|
||
conv.insert_many(docs)
|
||
|
||
# 3. 复合索引：遵循 ESR（等值 Equality → 排序 Sort → 范围 Range）
|
||
conv.create_index([("tenant_id", ASCENDING), ("status", ASCENDING), ("created_at", DESCENDING)], name="tenant_status_time")
|
||
query = {"tenant_id": "t1", "status": "handoff"}
|
||
plan = conv.find(query).sort("created_at", DESCENDING).limit(10).explain()["queryPlanner"]["winningPlan"]
|
||
stages = plan_stages(plan)
|
||
assert "IXSCAN" in stages and "COLLSCAN" not in stages and "SORT" not in stages, stages
|
||
print(f"2. 查询走索引且无需内存排序：{sorted(s for s in stages if s)}")
|
||
|
||
# 4. 聚合：每个租户每个渠道的会话数和平均 token
|
||
pipeline = [
|
||
{"$match": {"created_at": {"$gte": now - timedelta(hours=24)}}},
|
||
{"$unwind": "$messages"},
|
||
{"$match": {"messages.role": "assistant"}},
|
||
{"$group": {"_id": {"tenant": "$tenant_id", "channel": "$channel"}, "conversations": {"$sum": 1}, "avg_tokens": {"$avg": "$messages.tokens"}}},
|
||
{"$sort": {"_id.tenant": 1, "_id.channel": 1}},
|
||
]
|
||
rows = list(conv.aggregate(pipeline))
|
||
assert sum(r["conversations"] for r in rows) == 300
|
||
print(f"3. 聚合统计 {len(rows)} 组，例如 {rows[0]['_id']} 共 {rows[0]['conversations']} 个会话，平均 {rows[0]['avg_tokens']:.1f} tokens")
|
||
|
||
# 5. 幂等写入：外部消息 id 做唯一索引 + upsert，重复投递不会产生重复数据
|
||
events = db.events
|
||
events.create_index("event_id", unique=True)
|
||
for _ in range(3): # 同一条消息投递 3 次
|
||
events.update_one({"event_id": "evt-001"}, {"$setOnInsert": {"payload": "退款申请", "at": now}}, upsert=True)
|
||
assert events.count_documents({"event_id": "evt-001"}) == 1
|
||
try:
|
||
events.insert_one({"event_id": "evt-001"})
|
||
raise AssertionError("唯一索引应该拒绝")
|
||
except DuplicateKeyError:
|
||
pass
|
||
print("4. 幂等 upsert：重复投递 3 次只有 1 条")
|
||
|
||
# 6. 多文档事务：转人工时同时更新会话和创建工单，要么都成功要么都失败
|
||
tickets = db.tickets
|
||
target = conv.find_one({"status": "open"})
|
||
assert target is not None
|
||
with client.start_session() as s:
|
||
with s.start_transaction():
|
||
conv.update_one({"_id": target["_id"]}, {"$set": {"status": "handoff"}}, session=s)
|
||
tickets.insert_one({"conversation_id": target["_id"], "tenant_id": target["tenant_id"], "state": "new"}, session=s)
|
||
assert conv.find_one({"_id": target["_id"]})["status"] == "handoff" and tickets.count_documents({}) == 1
|
||
try:
|
||
with client.start_session() as s:
|
||
with s.start_transaction():
|
||
tickets.insert_one({"conversation_id": "x", "state": "new"}, session=s)
|
||
raise RuntimeError("模拟中途失败")
|
||
except RuntimeError:
|
||
pass
|
||
assert tickets.count_documents({}) == 1
|
||
print("5. 事务：成功的提交生效，失败的整体回滚")
|
||
|
||
# 7. 变更流：数据同步到搜索引擎 / 缓存的常用手段（可以用 resume token 断点续传）
|
||
with tickets.watch([{"$match": {"operationType": "insert"}}]) as stream:
|
||
tickets.insert_one({"conversation_id": "y", "tenant_id": "t2", "state": "new"})
|
||
change = stream.try_next()
|
||
while change is None:
|
||
change = stream.try_next()
|
||
token = stream.resume_token
|
||
assert change["fullDocument"]["conversation_id"] == "y" and token is not None
|
||
print("6. 变更流收到插入事件，并拿到 resume token 用于断点续传")
|
||
print("全部检查通过")
|
||
finally:
|
||
client.drop_database(db.name)
|
||
client.close()
|