# -*- coding: UTF-8 -*-
'''
@Project :domainScanDemo
@File    :add_sensitive_words_table.py
@IDE     :PyCharm
@Author  :梦伴
@Date    :2026/4/12 00:03
@explain : 为数据库添加sensitive_words表
'''

import psycopg2
from loguru import logger
from app.config import config

def add_sensitive_words_table():
    """
    为数据库添加sensitive_words表
    """
    try:
        # 连接数据库
        conn = psycopg2.connect(
            host=config.DB_HOST,
            port=config.DB_PORT,
            database=config.DB_DATABASE,
            user=config.DB_USER,
            password=config.DB_PASSWORD
        )
        cur = conn.cursor()
        logger.info("数据库连接成功")
        
        # 创建sensitive_words表
        create_sensitive_words_table = """
        CREATE TABLE IF NOT EXISTS sensitive_words (
            id SERIAL PRIMARY KEY,
            word VARCHAR(255) UNIQUE NOT NULL
        )
        """
        cur.execute(create_sensitive_words_table)
        logger.info("创建sensitive_words表成功")
        
        # 添加字段注释
        cur.execute("COMMENT ON COLUMN sensitive_words.id IS '主键ID'")
        cur.execute("COMMENT ON COLUMN sensitive_words.word IS '敏感词'")
        logger.info("添加sensitive_words表字段注释成功")
        
        # 提交事务
        conn.commit()
        logger.info("sensitive_words表添加完成")
        
        # 关闭连接
        cur.close()
        conn.close()
        
    except Exception as e:
        logger.error(f"执行失败: {e}")

if __name__ == "__main__":
    add_sensitive_words_table()