570392aa6f
* [add]修改迁移文件新建空白表结构 * Add installation guide and environment setup Added installation instructions and environment requirements for MemoryBear. * [delete]删除api，web的readme.md。只保留唯一readme.md * Fix database connection example in README Update database connection configuration example in README.
88 lines
2.8 KiB
Python
88 lines
2.8 KiB
Python
"""init_db_table
|
|
|
|
Revision ID: bf36acfdffe3
|
|
Revises: 5de5ec651b01
|
|
Create Date: 2025-11-13 11:31:57.671799
|
|
|
|
"""
|
|
from typing import Sequence, Union
|
|
|
|
from alembic import op
|
|
import sqlalchemy as sa
|
|
|
|
|
|
# revision identifiers, used by Alembic.
|
|
revision: str = 'bf36acfdffe3'
|
|
down_revision: Union[str, None] = '5de5ec651b01'
|
|
branch_labels: Union[str, Sequence[str], None] = None
|
|
depends_on: Union[str, Sequence[str], None] = None
|
|
|
|
|
|
|
|
def upgrade():
|
|
# 创建或更新统计信息的函数
|
|
op.execute("""
|
|
CREATE OR REPLACE FUNCTION update_knowledge_stats()
|
|
RETURNS TRIGGER AS $$
|
|
BEGIN
|
|
-- 处理 documents 表的插入、更新或删除
|
|
IF TG_TABLE_NAME = 'documents' THEN
|
|
-- 更新 knowledges 表的 doc_num
|
|
UPDATE knowledges
|
|
SET doc_num = (
|
|
SELECT COUNT(*)
|
|
FROM documents
|
|
WHERE kb_id = knowledges.id AND status = 1
|
|
)
|
|
WHERE id = NEW.kb_id OR id = OLD.kb_id;
|
|
|
|
-- 更新 knowledges 表的 chunk_num
|
|
UPDATE knowledges
|
|
SET chunk_num = (
|
|
SELECT COALESCE(SUM(chunk_num), 0)
|
|
FROM documents
|
|
WHERE kb_id = knowledges.id AND status = 1
|
|
)
|
|
WHERE id = NEW.kb_id OR id = OLD.kb_id;
|
|
|
|
-- 通过 knowledge_shares 表同步统计信息
|
|
-- 1. 使用 source_kb_id 的 doc_num 更新 target_kb_id 的 doc_num
|
|
UPDATE knowledges AS target
|
|
SET doc_num = source.doc_num
|
|
FROM knowledge_shares ks
|
|
JOIN knowledges AS source ON source.id = ks.source_kb_id
|
|
WHERE ks.target_kb_id = target.id
|
|
AND (source.id = NEW.kb_id OR source.id = OLD.kb_id);
|
|
|
|
-- 2. 使用 source_kb_id 的 chunk_num 更新 target_kb_id 的 chunk_num
|
|
UPDATE knowledges AS target
|
|
SET chunk_num = source.chunk_num
|
|
FROM knowledge_shares ks
|
|
JOIN knowledges AS source ON source.id = ks.source_kb_id
|
|
WHERE ks.target_kb_id = target.id
|
|
AND (source.id = NEW.kb_id OR source.id = OLD.kb_id);
|
|
|
|
END IF;
|
|
|
|
RETURN NULL;
|
|
END;
|
|
$$ LANGUAGE plpgsql;
|
|
""")
|
|
|
|
# 创建 documents 表上的触发器（如果不存在）
|
|
op.execute("""
|
|
DROP TRIGGER IF EXISTS tr_documents_update_stats ON documents;
|
|
CREATE TRIGGER tr_documents_update_stats
|
|
AFTER INSERT OR UPDATE OR DELETE ON documents
|
|
FOR EACH ROW
|
|
EXECUTE FUNCTION update_knowledge_stats();
|
|
""")
|
|
|
|
|
|
def downgrade():
|
|
# 删除触发器
|
|
op.execute("DROP TRIGGER IF EXISTS tr_documents_update_stats ON documents;")
|
|
|
|
# 删除函数
|
|
op.execute("DROP FUNCTION IF EXISTS update_knowledge_stats();")
|