RAG/src/main/java/com/wok/supportbot/config/DatabaseInitConfig.java

package com.wok.supportbot.config;

import jakarta.annotation.PostConstruct;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Component;

/**
 * 数据库初始化配置
 * 应用启动时检查并创建必要的表
 */
@Component
@Slf4j
public class DatabaseInitConfig {

    @Autowired
    private JdbcTemplate jdbcTemplate;

    @PostConstruct
    public void init() {
        try {
            // 检查 knowledge_category 表是否存在
            boolean categoryTableExists = checkTableExists("knowledge_category");
            if (!categoryTableExists) {
                log.info("创建知识库分类表 knowledge_category");
                createCategoryTable();
            }

            // 检查 knowledge_document 表是否存在
            boolean documentTableExists = checkTableExists("knowledge_document");
            if (!documentTableExists) {
                log.info("创建知识文档表 knowledge_document");
                createDocumentTable();
            } else {
                // 修复已存在表的 tags 默认值（从数组改为对象）
                fixTagsDefaultValue();
                // 自动添加 content_hash 列（二期新增）
                addContentHashColumn();
            }

            log.info("数据库初始化完成");
        } catch (Exception e) {
            log.error("数据库初始化失败", e);
        }
    }

    private boolean checkTableExists(String tableName) {
        try {
            String sql = "SELECT 1 FROM " + tableName + " LIMIT 1";
            jdbcTemplate.queryForObject(sql, Integer.class);
            return true;
        } catch (Exception e) {
            return false;
        }
    }

    private void createCategoryTable() {
        String sql = """
            CREATE TABLE IF NOT EXISTS knowledge_category (
                id              BIGSERIAL       PRIMARY KEY,
                name            VARCHAR(100)    NOT NULL,
                description     TEXT,
                parent_id       BIGINT          DEFAULT 0             NOT NULL,
                sort_order      INTEGER         DEFAULT 0             NOT NULL,
                document_count  INTEGER         DEFAULT 0             NOT NULL,
                create_time     TIMESTAMP       DEFAULT CURRENT_TIMESTAMP NOT NULL,
                is_delete       BOOLEAN         DEFAULT FALSE             NOT NULL
            )
            """;
        jdbcTemplate.execute(sql);

        // 创建索引
        jdbcTemplate.execute("CREATE INDEX IF NOT EXISTS idx_knowledge_category_parent ON knowledge_category (parent_id)");
    }

    private void createDocumentTable() {
        String sql = """
            CREATE TABLE IF NOT EXISTS knowledge_document (
                id              BIGSERIAL       PRIMARY KEY,
                title           VARCHAR(500)    NOT NULL,
                source_name     VARCHAR(500),
                file_type       VARCHAR(20)     NOT NULL,
                file_size       BIGINT          DEFAULT 0             NOT NULL,
                content         TEXT,
                category_id     BIGINT          DEFAULT 0             NOT NULL,
                tags            JSONB           DEFAULT '{}'          NOT NULL,
                chunk_count     INTEGER         DEFAULT 0             NOT NULL,
                status          VARCHAR(20)     DEFAULT 'PROCESSING'  NOT NULL,
                error_message   TEXT,
                create_time     TIMESTAMP       DEFAULT CURRENT_TIMESTAMP NOT NULL,
                update_time     TIMESTAMP       DEFAULT CURRENT_TIMESTAMP NOT NULL,
                is_delete       BOOLEAN         DEFAULT FALSE             NOT NULL
            )
            """;
        jdbcTemplate.execute(sql);

        // 创建索引
        jdbcTemplate.execute("CREATE INDEX IF NOT EXISTS idx_knowledge_document_category ON knowledge_document (category_id)");
        jdbcTemplate.execute("CREATE INDEX IF NOT EXISTS idx_knowledge_document_status ON knowledge_document (status)");
        jdbcTemplate.execute("CREATE INDEX IF NOT EXISTS idx_knowledge_document_create_time ON knowledge_document (create_time DESC)");
    }

    private void fixTagsDefaultValue() {
        try {
            // 检查当前默认值是否为数组
            String checkSql = "SELECT column_default FROM information_schema.columns WHERE table_name = 'knowledge_document' AND column_name = 'tags'";
            String currentDefault = jdbcTemplate.queryForObject(checkSql, String.class);
            if (currentDefault != null && currentDefault.contains("[]")) {
                log.info("修复 knowledge_document.tags 默认值");
                jdbcTemplate.execute("ALTER TABLE knowledge_document ALTER COLUMN tags SET DEFAULT '{}'");
                // 将已有的 '[]' 更新为 '{}'
                jdbcTemplate.execute("UPDATE knowledge_document SET tags = '{}' WHERE tags = '[]' OR tags IS NULL");
            }
        } catch (Exception e) {
            log.warn("修复 tags 默认值时出错（可能已修复）", e);
        }
    }

    /**
     * 自动添加 content_hash 列（二期去重功能新增字段）
     */
    private void addContentHashColumn() {
        try {
            String checkSql = "SELECT COUNT(*) FROM information_schema.columns WHERE table_name = 'knowledge_document' AND column_name = 'content_hash'";
            Integer count = jdbcTemplate.queryForObject(checkSql, Integer.class);
            if (count != null && count == 0) {
                log.info("添加 knowledge_document.content_hash 列");
                jdbcTemplate.execute("ALTER TABLE knowledge_document ADD COLUMN content_hash VARCHAR(64)");
                jdbcTemplate.execute("CREATE INDEX IF NOT EXISTS idx_knowledge_document_content_hash ON knowledge_document (content_hash)");
            }
        } catch (Exception e) {
            log.warn("添加 content_hash 列时出错", e);
        }
    }
}