first commit
This commit is contained in:
BF
263
init_database.py
Normal file
263
init_database.py
Normal file
@@ -0,0 +1,263 @@
|
||||
# -*- 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()
|
||||
Reference in New Issue
Block a user