# -*- coding: UTF-8 -*- ''' @Project :domainScanDemo @File :init_database.py @IDE :PyCharm @Author :梦伴 @Date :2026/4/9 21:20 @explain : 数据库初始化脚本 ''' import psycopg2 from loguru import logger from app.config import config def init_database(): """ 初始化数据库表结构 """ 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("数据库连接成功") # 创建domains表 create_domains_table = """ CREATE TABLE IF NOT EXISTS domains ( id SERIAL PRIMARY KEY, domain VARCHAR(255) UNIQUE NOT NULL, tld VARCHAR(50) NOT NULL, source_type INTEGER NOT NULL, use_status INTEGER DEFAULT 0, detect_status INTEGER DEFAULT 0, register_status INTEGER DEFAULT 0, has_beian BOOLEAN DEFAULT FALSE, company_type VARCHAR(100), website_url VARCHAR(255), beian_year INTEGER, snapshot_years TEXT, expire_date TIMESTAMP, create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP, update_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """ cur.execute(create_domains_table) logger.info("创建domains表成功") # 添加新字段 try: # 添加review_status字段 cur.execute("ALTER TABLE domains ADD COLUMN IF NOT EXISTS review_status INTEGER DEFAULT 0") # 添加detect_time字段 cur.execute("ALTER TABLE domains ADD COLUMN IF NOT EXISTS detect_time TIMESTAMP") # 添加backlink_count字段 cur.execute("ALTER TABLE domains ADD COLUMN IF NOT EXISTS backlink_count INTEGER DEFAULT 0") cur.execute("ALTER TABLE domains ADD COLUMN IF NOT EXISTS jucha_status INTEGER DEFAULT 0") cur.execute("ALTER TABLE domains ADD COLUMN IF NOT EXISTS juziseo_status INTEGER DEFAULT 0") # 尝试修改has_beian字段类型 try: # 先检查字段类型 cur.execute("SELECT data_type FROM information_schema.columns WHERE table_name = 'domains' AND column_name = 'has_beian'") result = cur.fetchone() if result and result[0] != 'integer': # 先删除默认值 cur.execute("ALTER TABLE domains ALTER COLUMN has_beian DROP DEFAULT") # 修改字段类型 cur.execute("ALTER TABLE domains ALTER COLUMN has_beian TYPE INTEGER USING CASE WHEN has_beian THEN 2 ELSE 3 END") # 设置新的默认值 cur.execute("ALTER TABLE domains ALTER COLUMN has_beian SET DEFAULT 1") except Exception as e: logger.warning(f"修改has_beian字段类型失败: {e}") logger.info("添加新字段成功") except Exception as e: logger.warning(f"添加新字段失败: {e}") # 回滚当前事务 conn.rollback() # 只添加新字段，不修改has_beian字段类型 try: # 添加review_status字段 cur.execute("ALTER TABLE domains ADD COLUMN IF NOT EXISTS review_status INTEGER DEFAULT 0") # 添加detect_time字段 cur.execute("ALTER TABLE domains ADD COLUMN IF NOT EXISTS detect_time TIMESTAMP") # 添加backlink_count字段 cur.execute("ALTER TABLE domains ADD COLUMN IF NOT EXISTS backlink_count INTEGER DEFAULT 0") cur.execute("ALTER TABLE domains ADD COLUMN IF NOT EXISTS jucha_status INTEGER DEFAULT 0") cur.execute("ALTER TABLE domains ADD COLUMN IF NOT EXISTS juziseo_status INTEGER DEFAULT 0") logger.info("添加部分新字段成功") except Exception as e2: logger.warning(f"添加新字段失败: {e2}") # 添加字段注释 cur.execute("COMMENT ON COLUMN domains.id IS '主键ID'") cur.execute("COMMENT ON COLUMN domains.domain IS '域名'") cur.execute("COMMENT ON COLUMN domains.tld IS '顶级域名'") cur.execute("COMMENT ON COLUMN domains.source_type IS '来源类型'") cur.execute("COMMENT ON COLUMN domains.use_status IS '使用状态'") cur.execute("COMMENT ON COLUMN domains.detect_status IS '检测状态'") cur.execute("COMMENT ON COLUMN domains.register_status IS '注册状态'") cur.execute("COMMENT ON COLUMN domains.review_status IS '人工复核状态'") cur.execute("COMMENT ON COLUMN domains.has_beian IS '是否有备案历史'") cur.execute("COMMENT ON COLUMN domains.company_type IS '单位性质'") cur.execute("COMMENT ON COLUMN domains.website_url IS '网站首页网址'") cur.execute("COMMENT ON COLUMN domains.beian_year IS '备案年份'") cur.execute("COMMENT ON COLUMN domains.snapshot_years IS '快照年份'") cur.execute("COMMENT ON COLUMN domains.expire_date IS '过期时间'") cur.execute("COMMENT ON COLUMN domains.detect_time IS '检测时间'") cur.execute("COMMENT ON COLUMN domains.jucha_status IS '聚查检测状态：0未检测，1已检测'") cur.execute("COMMENT ON COLUMN domains.juziseo_status IS '桔子检测状态：0未检测，1已检测'") cur.execute("COMMENT ON COLUMN domains.create_time IS '创建时间'") cur.execute("COMMENT ON COLUMN domains.update_time IS '更新时间'") logger.info("添加domains表字段注释成功") # 创建detect_tasks表 create_tasks_table = """ CREATE TABLE IF NOT EXISTS detect_tasks ( id SERIAL PRIMARY KEY, domain_id INTEGER REFERENCES domains(id), task_type INTEGER NOT NULL, status INTEGER DEFAULT 0, priority INTEGER DEFAULT 0, retry_count INTEGER DEFAULT 0, create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP, update_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """ cur.execute(create_tasks_table) logger.info("创建detect_tasks表成功") # 添加字段注释 cur.execute("COMMENT ON COLUMN detect_tasks.id IS '主键ID'") cur.execute("COMMENT ON COLUMN detect_tasks.domain_id IS '域名ID'") cur.execute("COMMENT ON COLUMN detect_tasks.task_type IS '任务类型'") cur.execute("COMMENT ON COLUMN detect_tasks.status IS '任务状态'") cur.execute("COMMENT ON COLUMN detect_tasks.priority IS '优先级'") cur.execute("COMMENT ON COLUMN detect_tasks.retry_count IS '重试次数'") cur.execute("COMMENT ON COLUMN detect_tasks.create_time IS '创建时间'") cur.execute("COMMENT ON COLUMN detect_tasks.update_time IS '更新时间'") logger.info("添加detect_tasks表字段注释成功") # 创建domain_blacklist表 create_blacklist_table = """ CREATE TABLE IF NOT EXISTS domain_blacklist ( id SERIAL PRIMARY KEY, domain VARCHAR(255) UNIQUE NOT NULL, reason TEXT, create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """ cur.execute(create_blacklist_table) logger.info("创建domain_blacklist表成功") # 添加字段注释 cur.execute("COMMENT ON COLUMN domain_blacklist.id IS '主键ID'") cur.execute("COMMENT ON COLUMN domain_blacklist.domain IS '域名'") cur.execute("COMMENT ON COLUMN domain_blacklist.reason IS '黑名单原因'") cur.execute("COMMENT ON COLUMN domain_blacklist.create_time IS '创建时间'") logger.info("添加domain_blacklist表字段注释成功") # 创建sensitive_words表 create_sensitive_words_table = """ CREATE TABLE IF NOT EXISTS sensitive_words ( id SERIAL PRIMARY KEY, word VARCHAR(255) UNIQUE NOT NULL, create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """ cur.execute(create_sensitive_words_table) logger.info("创建sensitive_words表成功") # 添加缺失的字段 try: # 添加category字段 cur.execute("ALTER TABLE sensitive_words ADD COLUMN IF NOT EXISTS category VARCHAR(50) DEFAULT 'default'") # 添加priority字段 cur.execute("ALTER TABLE sensitive_words ADD COLUMN IF NOT EXISTS priority INTEGER DEFAULT 1") logger.info("添加sensitive_words表字段成功") except Exception as e: logger.warning(f"添加sensitive_words表字段失败: {e}") # 添加字段注释 try: cur.execute("COMMENT ON COLUMN sensitive_words.id IS '主键ID'") cur.execute("COMMENT ON COLUMN sensitive_words.word IS '敏感词'") cur.execute("COMMENT ON COLUMN sensitive_words.category IS '分类'") cur.execute("COMMENT ON COLUMN sensitive_words.priority IS '优先级'") cur.execute("COMMENT ON COLUMN sensitive_words.create_time IS '创建时间'") logger.info("添加sensitive_words表字段注释成功") except Exception as e: logger.warning(f"添加sensitive_words表字段注释失败: {e}") # 创建domain_detections表 create_detections_table = """ CREATE TABLE IF NOT EXISTS domain_detections ( id SERIAL PRIMARY KEY, domain_id INTEGER REFERENCES domains(id), baidu_history JSONB, baidu_site JSONB, qihu360_site JSONB, google_site JSONB, wayback_info JSONB, chinaz_info JSONB, aizhan_info JSONB, juziseo_info JSONB, jucha_info JSONB, same_url BOOLEAN, is_chinese_title BOOLEAN, backlink_count_gt_10 BOOLEAN, create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP, update_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """ cur.execute(create_detections_table) logger.info("创建domain_detections表成功") cur.execute("CREATE UNIQUE INDEX IF NOT EXISTS idx_domain_detections_domain_id ON domain_detections(domain_id)") # 为已存在的domain_detections表添加新字段 try: # 添加same_url字段 cur.execute("ALTER TABLE domain_detections ADD COLUMN IF NOT EXISTS same_url BOOLEAN") # 添加is_chinese_title字段 cur.execute("ALTER TABLE domain_detections ADD COLUMN IF NOT EXISTS is_chinese_title BOOLEAN") # 添加backlink_count_gt_10字段 cur.execute("ALTER TABLE domain_detections ADD COLUMN IF NOT EXISTS backlink_count_gt_10 BOOLEAN") # 添加wayback_info字段 cur.execute("ALTER TABLE domain_detections ADD COLUMN IF NOT EXISTS wayback_info JSONB") logger.info("为domain_detections表添加新字段成功") except Exception as e: logger.warning(f"为domain_detections表添加新字段失败: {e}") # 回滚当前事务 conn.rollback() # 添加字段注释 cur.execute("COMMENT ON COLUMN domain_detections.id IS '主键ID'") cur.execute("COMMENT ON COLUMN domain_detections.domain_id IS '域名ID'") cur.execute("COMMENT ON COLUMN domain_detections.baidu_history IS '百度历史'") cur.execute("COMMENT ON COLUMN domain_detections.baidu_site IS '百度site'") cur.execute("COMMENT ON COLUMN domain_detections.qihu360_site IS '360 site'") cur.execute("COMMENT ON COLUMN domain_detections.google_site IS 'Google site'") cur.execute("COMMENT ON COLUMN domain_detections.wayback_info IS '时光机检测信息'") cur.execute("COMMENT ON COLUMN domain_detections.chinaz_info IS '站长之家信息'") cur.execute("COMMENT ON COLUMN domain_detections.aizhan_info IS '爱站网信息'") cur.execute("COMMENT ON COLUMN domain_detections.juziseo_info IS '桔子SEO信息'") cur.execute("COMMENT ON COLUMN domain_detections.jucha_info IS '聚查信息'") cur.execute("COMMENT ON COLUMN domain_detections.same_url IS '首网址和备案网址是否一样'") cur.execute("COMMENT ON COLUMN domain_detections.is_chinese_title IS 'title是否简体中文'") cur.execute("COMMENT ON COLUMN domain_detections.backlink_count_gt_10 IS '快照历史友链数量是否大于10'") cur.execute("COMMENT ON COLUMN domain_detections.create_time IS '创建时间'") cur.execute("COMMENT ON COLUMN domain_detections.update_time IS '更新时间'") logger.info("添加domain_detections表字段注释成功") # 提交事务 conn.commit() logger.info("数据库初始化完成") except Exception as e: logger.error(f"数据库初始化失败: {e}") if conn: conn.rollback() finally: if cur: cur.close() if conn: conn.close() if __name__ == "__main__": init_database()