# -*- coding: UTF-8 -*-
'''
@Project :domainScanDemo
@File    :add_expire_date_field.py
@IDE     :PyCharm
@Author  :梦伴
@Date    :2026/4/11 22:40
@explain : 为domains表添加expire_date字段
'''

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

def add_expire_date_field():
    """
    为domains表添加expire_date字段
    """
    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("数据库连接成功")
        
        # 添加expire_date字段
        try:
            cur.execute("ALTER TABLE domains ADD COLUMN IF NOT EXISTS expire_date DATE")
            cur.execute("COMMENT ON COLUMN domains.expire_date IS '过期时间'")
            logger.info("添加expire_date字段成功")
        except Exception as e:
            logger.error(f"添加expire_date字段失败: {e}")
        
        # 提交事务
        conn.commit()
        logger.info("字段添加完成")
        
        # 关闭连接
        cur.close()
        conn.close()
        
    except Exception as e:
        logger.error(f"执行失败: {e}")

if __name__ == "__main__":
    add_expire_date_field()