263 lines
12 KiB
Python
263 lines
12 KiB
Python
# -*- 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")
|
||
# 尝试修改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")
|
||
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.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,
|
||
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表成功")
|
||
|
||
# 为已存在的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")
|
||
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.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() |