Agent/knowhub/knowhub_db/cascade.py

"""
级联删除：Greenplum 不支持 FK ON DELETE CASCADE，由应用层保证。

用法：
    from knowhub.knowhub_db.cascade import cascade_delete
    cascade_delete(cursor, 'knowledge', knowledge_id)
"""

# 每个实体表涉及的关联表及其外键列名
_JUNCTIONS = {
    'knowledge': [
        ('requirement_knowledge', 'knowledge_id'),
        ('capability_knowledge', 'knowledge_id'),
        ('tool_knowledge', 'knowledge_id'),
        ('knowledge_resource', 'knowledge_id'),
        ('knowledge_relation', 'source_id'),
        ('knowledge_relation', 'target_id'),
        ('strategy_knowledge', 'knowledge_id'),
    ],
    'tool': [
        ('capability_tool', 'tool_id'),
        ('tool_knowledge', 'tool_id'),
        ('tool_provider', 'tool_id'),
    ],
    'capability': [
        ('requirement_capability', 'capability_id'),
        ('capability_tool', 'capability_id'),
        ('capability_knowledge', 'capability_id'),
        ('capability_resource', 'capability_id'),
        ('strategy_capability', 'capability_id'),
    ],
    'requirement': [
        ('requirement_capability', 'requirement_id'),
        ('requirement_knowledge', 'requirement_id'),
        ('requirement_resource', 'requirement_id'),
        ('requirement_strategy', 'requirement_id'),
    ],
    'resource': [
        ('knowledge_resource', 'resource_id'),
        ('capability_resource', 'resource_id'),
        ('requirement_resource', 'resource_id'),
        ('strategy_resource', 'resource_id'),
    ],
    'strategy': [
        ('strategy_capability', 'strategy_id'),
        ('strategy_knowledge', 'strategy_id'),
        ('strategy_resource', 'strategy_id'),
        ('requirement_strategy', 'strategy_id'),
    ],
}


def cascade_delete(cursor, entity_table: str, entity_id: str):
    """先删除关联表中的引用行，再删除实体本身"""
    for junction_table, fk_column in _JUNCTIONS.get(entity_table, []):
        cursor.execute(
            f"DELETE FROM {junction_table} WHERE {fk_column} = %s",
            (entity_id,))
    cursor.execute(
        f"DELETE FROM {entity_table} WHERE id = %s",
        (entity_id,))


# 带 version 字段的实体表（tool 不在其中——它的 version 是工具自身发布版本，不是多租户标签）
_VERSIONED_ENTITIES = ['knowledge', 'resource', 'requirement', 'capability', 'strategy']


def purge_version(cursor, version: str) -> dict:
    """
    批量清除指定 version 的所有实体（及其 junction）。
    常用于测试数据迭代——整批抹掉 tao_dev_1，重新入库。

    Returns: {table: row_count_deleted, ...}
    """
    if not version or version == 'v0':
        raise ValueError(f"Refusing to purge protected version: {version!r}（v0 和空字符串受保护）")

    stats = {}
    for table in _VERSIONED_ENTITIES:
        # 先清掉所有指向这个版本实体的 junction 行
        for junction, fk_col in _JUNCTIONS.get(table, []):
            cursor.execute(
                f"DELETE FROM {junction} WHERE {fk_col} IN "
                f"(SELECT id FROM {table} WHERE version = %s)",
                (version,))
        # 再删实体本身
        cursor.execute(f"DELETE FROM {table} WHERE version = %s", (version,))
        stats[table] = cursor.rowcount if cursor.rowcount is not None else 0
    return stats