2743 lines
102 KiB
Python
2743 lines
102 KiB
Python
# -*- coding: UTF-8 -*-
|
||
'''
|
||
@Project :domainScanDemo
|
||
@File :database.py
|
||
@IDE :PyCharm
|
||
@Author :梦伴
|
||
@Date :2026/4/9 0:07
|
||
@explain : 数据库操作类
|
||
'''
|
||
|
||
import json
|
||
import os
|
||
import socket
|
||
import threading
|
||
import time
|
||
|
||
import psycopg2
|
||
import redis
|
||
from loguru import logger
|
||
from psycopg2 import extensions
|
||
from psycopg2.extras import Json
|
||
from app.config import config
|
||
from app.utils.status_codes import (
|
||
DETECT_STATUS_BLACKLISTED,
|
||
DETECT_STATUS_COMPLETED,
|
||
DETECT_STATUS_FAILED,
|
||
DETECT_STATUS_PENDING,
|
||
DETECT_STATUS_RUNNING,
|
||
REGISTER_STATUS_AVAILABLE,
|
||
REGISTER_STATUS_REGISTERED,
|
||
REVIEW_STATUS_PENDING,
|
||
THIRD_PARTY_STATUS_DONE,
|
||
)
|
||
|
||
|
||
_POOL_VERBOSE_LOGS = bool(int(os.getenv("DOMAINCHECK_DB_POOL_VERBOSE_LOGS", "0") or 0))
|
||
|
||
|
||
_STEP_CLAIM_PRIORITY = {
|
||
# Prefer deeper pipeline steps first so domains that already passed an
|
||
# earlier gate can keep advancing instead of being starved behind the
|
||
# oldest first-step backlog. This keeps later steps like 360/chinaz/aizhan
|
||
# from sitting pending forever while baidu/register continue to dominate
|
||
# the queue.
|
||
"detect_juziseo": 10,
|
||
"detect_jucha": 20,
|
||
"detect_wayback": 30,
|
||
"detect_aizhan": 40,
|
||
"detect_chinaz": 50,
|
||
"detect_360_site": 60,
|
||
"detect_baidu_site": 70,
|
||
"detect_register": 100,
|
||
}
|
||
|
||
|
||
def _detect_job_item_step_priority(step_code: str) -> int:
|
||
normalized = str(step_code or "").strip()
|
||
if not normalized:
|
||
return 999
|
||
return int(_STEP_CLAIM_PRIORITY.get(normalized, 10))
|
||
|
||
|
||
def _ordered_step_claim_codes() -> list[str]:
|
||
return [
|
||
step_code
|
||
for step_code, _ in sorted(
|
||
_STEP_CLAIM_PRIORITY.items(),
|
||
key=lambda item: (int(item[1] or 0), str(item[0] or "")),
|
||
)
|
||
]
|
||
|
||
|
||
def _resolve_step_claim_quota(limit: int) -> int:
|
||
configured_quota = int(os.getenv("DOMAINCHECK_STEP_CLAIM_QUOTA", "0") or 0)
|
||
if configured_quota > 0:
|
||
return max(1, configured_quota)
|
||
normalized_limit = max(1, int(limit or 1))
|
||
return max(64, min(normalized_limit, max(1, normalized_limit // 4)))
|
||
|
||
|
||
class Database:
|
||
"""
|
||
数据库操作类
|
||
"""
|
||
|
||
def __init__(self, host=None, port=None, database=None, user=None, password=None):
|
||
"""
|
||
初始化数据库连接
|
||
|
||
:param host: 数据库主机
|
||
:param port: 数据库端口
|
||
:param database: 数据库名称
|
||
:param user: 用户名
|
||
:param password: 密码
|
||
"""
|
||
self.host = host or config.DB_HOST
|
||
self.port = port or config.DB_PORT
|
||
self.database = database or config.DB_DATABASE
|
||
self.user = user or config.DB_USER
|
||
self.password = password or config.DB_PASSWORD
|
||
|
||
# 数据库连接池
|
||
self.connection_pool = []
|
||
self.pool_size = max(1, int(config.DB_POOL_SIZE)) # 连接池大小
|
||
self.pool_warm_size = max(1, min(self.pool_size, int(getattr(config, 'DB_POOL_WARM_SIZE', 16) or 16)))
|
||
self.pool_idle_keep_max = max(
|
||
self.pool_warm_size,
|
||
min(self.pool_size, int(getattr(config, 'DB_POOL_IDLE_KEEP_MAX', self.pool_size) or self.pool_size)),
|
||
)
|
||
self.pool_healthcheck_interval = max(
|
||
0.0,
|
||
float(getattr(config, 'DB_POOL_HEALTHCHECK_INTERVAL', 30) or 30),
|
||
)
|
||
self.pool_acquire_timeout = max(1.0, float(getattr(config, 'DB_POOL_ACQUIRE_TIMEOUT', 20)))
|
||
self.pool_lock = threading.Lock()
|
||
self.pool_condition = threading.Condition(self.pool_lock)
|
||
self._pool_initialized = False
|
||
self.total_connections = 0
|
||
self._connection_last_healthcheck = {}
|
||
|
||
# 初始化连接池
|
||
self._init_connection_pool()
|
||
|
||
# 初始化 Redis 客户端
|
||
try:
|
||
self.redis_client = redis.Redis(
|
||
host=config.REDIS_HOST,
|
||
port=config.REDIS_PORT,
|
||
password=config.REDIS_PASSWORD,
|
||
db=config.REDIS_DB,
|
||
decode_responses=True
|
||
)
|
||
# 测试连接
|
||
self.redis_client.ping()
|
||
logger.info(f"Redis 连接成功: {config.REDIS_HOST}:{config.REDIS_PORT}")
|
||
self.use_redis = True
|
||
|
||
# 初始化布隆过滤器
|
||
self._init_bloom_filter()
|
||
except Exception as e:
|
||
logger.warning(f"Redis 连接失败: {e}，将使用数据库查询")
|
||
self.redis_client = None
|
||
self.use_redis = False
|
||
self.use_bloom_filter = False
|
||
|
||
def _init_connection_pool(self):
|
||
"""
|
||
初始化数据库连接池
|
||
"""
|
||
try:
|
||
with self.pool_condition:
|
||
missing = self.pool_warm_size - self.total_connections
|
||
if missing <= 0:
|
||
return
|
||
|
||
created = []
|
||
for i in range(missing):
|
||
created.append(self._create_connection())
|
||
|
||
with self.pool_condition:
|
||
self.connection_pool.extend(created)
|
||
self.total_connections += len(created)
|
||
self._pool_initialized = True
|
||
now = time.monotonic()
|
||
for conn in created:
|
||
self._connection_last_healthcheck[id(conn)] = now
|
||
self.pool_condition.notify_all()
|
||
logger.info(
|
||
f"数据库连接池初始化成功，预热: {len(created)}，池中空闲: {len(self.connection_pool)}，总连接: {self.total_connections}/{self.pool_size}"
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"初始化数据库连接池失败: {e}")
|
||
|
||
def _create_connection(self):
|
||
return psycopg2.connect(
|
||
host=self.host,
|
||
port=self.port,
|
||
database=self.database,
|
||
user=self.user,
|
||
password=self.password,
|
||
connect_timeout=5,
|
||
application_name="domaincheck-worker",
|
||
)
|
||
|
||
def _init_bloom_filter(self):
|
||
"""
|
||
初始化布隆过滤器
|
||
"""
|
||
try:
|
||
# 检查 Redis 是否支持布隆过滤器
|
||
# 如果不支持，将使用普通缓存
|
||
try:
|
||
# 尝试创建布隆过滤器
|
||
self.redis_client.execute_command('BF.RESERVE', 'domain_bloom', 0.001, 1073741824)
|
||
logger.info("布隆过滤器初始化成功")
|
||
self.use_bloom_filter = True
|
||
except Exception as e:
|
||
# 检查是否是因为布隆过滤器已存在
|
||
if "item exists" in str(e):
|
||
logger.info("布隆过滤器已存在，直接使用")
|
||
self.use_bloom_filter = True
|
||
else:
|
||
# 如果命令不存在，说明 Redis 没有加载布隆过滤器模块
|
||
logger.info(f"Redis 布隆过滤器不可用: {e}，将使用普通缓存")
|
||
self.use_bloom_filter = False
|
||
except Exception as e:
|
||
logger.warning(f"初始化布隆过滤器失败: {e}")
|
||
self.use_bloom_filter = False
|
||
|
||
def connect(self, thread_id=None):
|
||
"""
|
||
从连接池获取数据库连接
|
||
|
||
:param thread_id: 线程ID，为None时使用当前线程ID
|
||
:return: tuple - (连接对象, 游标对象)
|
||
"""
|
||
import threading
|
||
thread_id = thread_id or threading.current_thread().ident
|
||
|
||
try:
|
||
if not self._pool_initialized:
|
||
self._init_connection_pool()
|
||
|
||
deadline = time.monotonic() + self.pool_acquire_timeout
|
||
warned_pool_empty = False
|
||
|
||
while True:
|
||
create_new = False
|
||
with self.pool_condition:
|
||
while self.connection_pool:
|
||
conn = self.connection_pool.pop()
|
||
if conn and not conn.closed:
|
||
try:
|
||
if conn.get_transaction_status() != extensions.TRANSACTION_STATUS_IDLE:
|
||
conn.rollback()
|
||
last_healthcheck = float(self._connection_last_healthcheck.get(id(conn), 0.0) or 0.0)
|
||
now = time.monotonic()
|
||
if self.pool_healthcheck_interval > 0 and now - last_healthcheck >= self.pool_healthcheck_interval:
|
||
cur = conn.cursor()
|
||
cur.execute("SELECT 1")
|
||
cur.fetchone()
|
||
cur.close()
|
||
self._connection_last_healthcheck[id(conn)] = now
|
||
if _POOL_VERBOSE_LOGS:
|
||
logger.debug(f"线程 {thread_id} 从连接池获取连接成功")
|
||
return conn, conn.cursor()
|
||
except Exception:
|
||
self._connection_last_healthcheck.pop(id(conn), None)
|
||
try:
|
||
conn.close()
|
||
except Exception:
|
||
pass
|
||
self.total_connections = max(0, self.total_connections - 1)
|
||
continue
|
||
|
||
if self.total_connections < self.pool_size:
|
||
self.total_connections += 1
|
||
create_new = True
|
||
else:
|
||
remaining = deadline - time.monotonic()
|
||
if remaining <= 0:
|
||
logger.warning(
|
||
f"连接池耗尽，线程 {thread_id} 等待超时 {self.pool_acquire_timeout}s，返回空连接"
|
||
)
|
||
return None, None
|
||
|
||
if not warned_pool_empty:
|
||
warned_pool_empty = True
|
||
logger.warning(
|
||
f"连接池耗尽，线程 {thread_id} 等待可复用连接，池大小 {self.pool_size}"
|
||
)
|
||
self.pool_condition.wait(timeout=min(0.5, remaining))
|
||
|
||
if create_new:
|
||
try:
|
||
conn = self._create_connection()
|
||
self._connection_last_healthcheck[id(conn)] = time.monotonic()
|
||
if _POOL_VERBOSE_LOGS:
|
||
logger.debug(f"线程 {thread_id} 新建数据库连接成功，总连接 {self.total_connections}/{self.pool_size}")
|
||
return conn, conn.cursor()
|
||
except Exception as e:
|
||
with self.pool_condition:
|
||
self.total_connections = max(0, self.total_connections - 1)
|
||
self.pool_condition.notify()
|
||
logger.error(f"线程 {thread_id} 新建数据库连接失败: {e}")
|
||
return None, None
|
||
except Exception as e:
|
||
logger.error(f"线程 {thread_id} 获取数据库连接失败: {e}")
|
||
return None, None
|
||
|
||
def get_connection(self):
|
||
"""
|
||
获取数据库连接（兼容方法）
|
||
|
||
:return: 连接对象
|
||
"""
|
||
conn, _ = self.connect()
|
||
return conn
|
||
|
||
def close(self, conn=None, cur=None):
|
||
"""
|
||
将数据库连接放回连接池
|
||
|
||
:param conn: 连接对象
|
||
:param cur: 游标对象
|
||
"""
|
||
try:
|
||
if cur:
|
||
try:
|
||
cur.close()
|
||
except:
|
||
pass
|
||
|
||
if conn and not conn.closed:
|
||
with self.pool_condition:
|
||
try:
|
||
if conn.get_transaction_status() != extensions.TRANSACTION_STATUS_IDLE:
|
||
conn.rollback()
|
||
except Exception:
|
||
try:
|
||
self._connection_last_healthcheck.pop(id(conn), None)
|
||
conn.close()
|
||
finally:
|
||
self.total_connections = max(0, self.total_connections - 1)
|
||
self.pool_condition.notify()
|
||
logger.warning("连接归还前回滚失败，已关闭连接")
|
||
return
|
||
if len(self.connection_pool) < self.pool_idle_keep_max:
|
||
self.connection_pool.append(conn)
|
||
self.pool_condition.notify()
|
||
if _POOL_VERBOSE_LOGS:
|
||
logger.debug("连接已放回连接池")
|
||
else:
|
||
# 启动阶段只预热少量连接，高并发跑起来后允许保留更多空闲连接，
|
||
# 避免本地数据库在“建连/关连”之间来回抖动。
|
||
self._connection_last_healthcheck.pop(id(conn), None)
|
||
conn.close()
|
||
self.total_connections = max(0, self.total_connections - 1)
|
||
self.pool_condition.notify()
|
||
if _POOL_VERBOSE_LOGS:
|
||
logger.debug(f"空闲连接超过保留阈值({self.pool_idle_keep_max})，已关闭多余连接")
|
||
except Exception as e:
|
||
logger.error(f"关闭数据库连接失败: {e}")
|
||
try:
|
||
if conn and not conn.closed:
|
||
self._connection_last_healthcheck.pop(id(conn), None)
|
||
conn.close()
|
||
except:
|
||
pass
|
||
|
||
def get_sensitive_words(self):
|
||
"""
|
||
获取所有敏感词
|
||
|
||
:return: list - 敏感词列表
|
||
"""
|
||
try:
|
||
sql = "SELECT word, category, priority FROM sensitive_words ORDER BY priority DESC, word ASC"
|
||
result = self.fetch_all(sql)
|
||
return result
|
||
except Exception as e:
|
||
logger.error(f"获取敏感词失败: {e}")
|
||
return []
|
||
|
||
def add_sensitive_word(self, word, category='default', priority=1):
|
||
"""
|
||
添加敏感词
|
||
|
||
:param word: 敏感词
|
||
:param category: 分类
|
||
:param priority: 优先级
|
||
:return: bool - 是否成功
|
||
"""
|
||
try:
|
||
sql = "INSERT INTO sensitive_words (word, category, priority) VALUES (%s, %s, %s) ON CONFLICT (word) DO NOTHING"
|
||
return self.execute(sql, (word, category, priority))
|
||
except Exception as e:
|
||
logger.error(f"添加敏感词失败: {word}, 错误: {e}")
|
||
return False
|
||
|
||
def delete_sensitive_word(self, word):
|
||
"""
|
||
删除敏感词
|
||
|
||
:param word: 敏感词
|
||
:return: bool - 是否成功
|
||
"""
|
||
try:
|
||
sql = "DELETE FROM sensitive_words WHERE word = %s"
|
||
return self.execute(sql, (word,))
|
||
except Exception as e:
|
||
logger.error(f"删除敏感词失败: {word}, 错误: {e}")
|
||
return False
|
||
|
||
def update_sensitive_word(self, old_word, new_word, category=None, priority=None):
|
||
"""
|
||
更新敏感词
|
||
|
||
:param old_word: 旧敏感词
|
||
:param new_word: 新敏感词
|
||
:param category: 分类
|
||
:param priority: 优先级
|
||
:return: bool - 是否成功
|
||
"""
|
||
try:
|
||
if category is not None and priority is not None:
|
||
sql = "UPDATE sensitive_words SET word = %s, category = %s, priority = %s WHERE word = %s"
|
||
return self.execute(sql, (new_word, category, priority, old_word))
|
||
elif category is not None:
|
||
sql = "UPDATE sensitive_words SET word = %s, category = %s WHERE word = %s"
|
||
return self.execute(sql, (new_word, category, old_word))
|
||
elif priority is not None:
|
||
sql = "UPDATE sensitive_words SET word = %s, priority = %s WHERE word = %s"
|
||
return self.execute(sql, (new_word, priority, old_word))
|
||
else:
|
||
sql = "UPDATE sensitive_words SET word = %s WHERE word = %s"
|
||
return self.execute(sql, (new_word, old_word))
|
||
except Exception as e:
|
||
logger.error(f"更新敏感词失败: {old_word} -> {new_word}, 错误: {e}")
|
||
return False
|
||
|
||
def batch_add_sensitive_words(self, words):
|
||
"""
|
||
批量添加敏感词
|
||
|
||
:param words: 敏感词列表，每个元素是 (word, category, priority) 元组
|
||
:return: bool - 是否成功
|
||
"""
|
||
try:
|
||
if not words:
|
||
return True
|
||
|
||
sql = "INSERT INTO sensitive_words (word, category, priority) VALUES (%s, %s, %s) ON CONFLICT (word) DO NOTHING"
|
||
return self.execute_many(sql, words)
|
||
except Exception as e:
|
||
logger.error(f"批量添加敏感词失败: {e}")
|
||
return False
|
||
|
||
def execute(self, sql, params=None):
|
||
"""
|
||
执行SQL语句
|
||
|
||
:param sql: SQL语句
|
||
:param params: 参数
|
||
:return: 执行结果
|
||
"""
|
||
import threading
|
||
thread_id = threading.current_thread().ident
|
||
conn = None
|
||
cur = None
|
||
|
||
try:
|
||
conn, cur = self.connect(thread_id)
|
||
if not conn or not cur:
|
||
logger.error(f"线程 {thread_id} 数据库连接失败，无法执行SQL")
|
||
return False
|
||
|
||
cur.execute(sql, params)
|
||
conn.commit()
|
||
return True
|
||
except Exception as e:
|
||
logger.error(f"执行SQL失败: {sql}, 错误: {e}")
|
||
try:
|
||
if conn:
|
||
conn.rollback()
|
||
except:
|
||
pass
|
||
return False
|
||
finally:
|
||
# 将连接放回连接池
|
||
self.close(conn, cur)
|
||
|
||
def execute_many(self, sql, params_list):
|
||
"""
|
||
批量执行SQL语句
|
||
|
||
:param sql: SQL语句
|
||
:param params_list: 参数列表
|
||
:return: 执行结果
|
||
"""
|
||
import threading
|
||
thread_id = threading.current_thread().ident
|
||
conn = None
|
||
cur = None
|
||
|
||
try:
|
||
conn, cur = self.connect(thread_id)
|
||
if not conn or not cur:
|
||
logger.error(f"线程 {thread_id} 数据库连接失败，无法执行批量SQL")
|
||
return False
|
||
|
||
cur.executemany(sql, params_list)
|
||
conn.commit()
|
||
return True
|
||
except Exception as e:
|
||
logger.error(f"执行批量SQL失败: {sql}, 错误: {e}")
|
||
try:
|
||
if conn:
|
||
conn.rollback()
|
||
except:
|
||
pass
|
||
return False
|
||
finally:
|
||
self.close(conn, cur)
|
||
|
||
def ensure_cluster_runtime_tables(self):
|
||
sql = """
|
||
CREATE TABLE IF NOT EXISTS detect_worker_nodes (
|
||
node_code VARCHAR(64) PRIMARY KEY,
|
||
region VARCHAR(32) NOT NULL DEFAULT 'unknown',
|
||
role VARCHAR(32) NOT NULL DEFAULT 'worker',
|
||
hostname VARCHAR(255) NOT NULL DEFAULT '',
|
||
ip VARCHAR(64) NOT NULL DEFAULT '',
|
||
status VARCHAR(32) NOT NULL DEFAULT 'unknown',
|
||
worker_version VARCHAR(32) NOT NULL DEFAULT '',
|
||
current_load INTEGER NOT NULL DEFAULT 0,
|
||
metadata_json JSONB,
|
||
last_heartbeat_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||
create_time TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||
update_time TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS detect_jobs (
|
||
id BIGSERIAL PRIMARY KEY,
|
||
job_code VARCHAR(64) NOT NULL UNIQUE,
|
||
source VARCHAR(64) NOT NULL DEFAULT 'manual',
|
||
plan_hash VARCHAR(128) NOT NULL DEFAULT '',
|
||
task_mode VARCHAR(32) NOT NULL DEFAULT 'domain_pipeline',
|
||
step_code VARCHAR(64) NOT NULL DEFAULT '',
|
||
status VARCHAR(32) NOT NULL DEFAULT 'pending',
|
||
remark TEXT NOT NULL DEFAULT '',
|
||
created_by VARCHAR(64) NOT NULL DEFAULT '',
|
||
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||
started_at TIMESTAMP,
|
||
finished_at TIMESTAMP
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS detect_job_items (
|
||
id BIGSERIAL PRIMARY KEY,
|
||
job_id BIGINT NOT NULL REFERENCES detect_jobs(id) ON DELETE CASCADE,
|
||
domain_id BIGINT NOT NULL,
|
||
step_code VARCHAR(64) NOT NULL DEFAULT '',
|
||
status VARCHAR(32) NOT NULL DEFAULT 'pending',
|
||
claimed_by VARCHAR(64) NOT NULL DEFAULT '',
|
||
claim_token VARCHAR(64) NOT NULL DEFAULT '',
|
||
lease_expires_at TIMESTAMP,
|
||
attempt_count INTEGER NOT NULL DEFAULT 0,
|
||
last_error TEXT NOT NULL DEFAULT '',
|
||
result_version VARCHAR(64) NOT NULL DEFAULT '',
|
||
step_payload_json JSONB,
|
||
result_payload_json JSONB,
|
||
started_at TIMESTAMP,
|
||
finished_at TIMESTAMP,
|
||
updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||
create_time TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||
);
|
||
|
||
CREATE INDEX IF NOT EXISTS idx_detect_job_items_status_lease
|
||
ON detect_job_items(status, lease_expires_at);
|
||
|
||
ALTER TABLE detect_jobs
|
||
ADD COLUMN IF NOT EXISTS task_mode VARCHAR(32) NOT NULL DEFAULT 'domain_pipeline',
|
||
ADD COLUMN IF NOT EXISTS step_code VARCHAR(64) NOT NULL DEFAULT '';
|
||
|
||
ALTER TABLE detect_job_items
|
||
ADD COLUMN IF NOT EXISTS step_code VARCHAR(64) NOT NULL DEFAULT '',
|
||
ADD COLUMN IF NOT EXISTS step_payload_json JSONB,
|
||
ADD COLUMN IF NOT EXISTS result_payload_json JSONB;
|
||
|
||
ALTER TABLE detect_job_items
|
||
DROP CONSTRAINT IF EXISTS uq_detect_job_items_job_domain;
|
||
|
||
CREATE UNIQUE INDEX IF NOT EXISTS idx_detect_job_items_job_domain_step
|
||
ON detect_job_items(job_id, domain_id, step_code);
|
||
|
||
CREATE TABLE IF NOT EXISTS detect_run_events (
|
||
id BIGSERIAL PRIMARY KEY,
|
||
job_id BIGINT REFERENCES detect_jobs(id) ON DELETE SET NULL,
|
||
job_item_id BIGINT REFERENCES detect_job_items(id) ON DELETE SET NULL,
|
||
node_code VARCHAR(64) NOT NULL DEFAULT '',
|
||
event_type VARCHAR(64) NOT NULL DEFAULT '',
|
||
level VARCHAR(16) NOT NULL DEFAULT 'info',
|
||
message TEXT NOT NULL DEFAULT '',
|
||
payload_json JSONB,
|
||
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS detect_sync_records (
|
||
id BIGSERIAL PRIMARY KEY,
|
||
sync_type VARCHAR(32) NOT NULL DEFAULT '',
|
||
source_region VARCHAR(32) NOT NULL DEFAULT '',
|
||
target_region VARCHAR(32) NOT NULL DEFAULT '',
|
||
status VARCHAR(32) NOT NULL DEFAULT 'pending',
|
||
payload_json JSONB,
|
||
error_message TEXT NOT NULL DEFAULT '',
|
||
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||
updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||
);
|
||
"""
|
||
return self.execute(sql)
|
||
|
||
def register_cluster_node(self, node_code, region, role, status='online', current_load=0, metadata=None):
|
||
try:
|
||
ip_addr = ''
|
||
try:
|
||
ip_addr = socket.gethostbyname(socket.gethostname())
|
||
except Exception:
|
||
pass
|
||
sql = """
|
||
INSERT INTO detect_worker_nodes (
|
||
node_code, region, role, hostname, ip, status, worker_version, current_load, metadata_json, last_heartbeat_at, update_time
|
||
) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s::jsonb, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
|
||
ON CONFLICT (node_code) DO UPDATE SET
|
||
region = EXCLUDED.region,
|
||
role = EXCLUDED.role,
|
||
hostname = EXCLUDED.hostname,
|
||
ip = EXCLUDED.ip,
|
||
status = EXCLUDED.status,
|
||
worker_version = EXCLUDED.worker_version,
|
||
current_load = EXCLUDED.current_load,
|
||
metadata_json = EXCLUDED.metadata_json,
|
||
last_heartbeat_at = CURRENT_TIMESTAMP,
|
||
update_time = CURRENT_TIMESTAMP
|
||
"""
|
||
return self.execute(
|
||
sql,
|
||
(
|
||
node_code,
|
||
region,
|
||
role,
|
||
socket.gethostname(),
|
||
ip_addr,
|
||
status,
|
||
'0.1.0',
|
||
int(current_load or 0),
|
||
json.dumps(metadata or {}, ensure_ascii=False),
|
||
),
|
||
)
|
||
except Exception as e:
|
||
logger.warning(f"注册检测节点失败: {e}")
|
||
return False
|
||
|
||
def claim_detect_job_items(self, node_code, limit=1000, lease_seconds=3600, job_id=None):
|
||
"""
|
||
领取一批待执行的任务项。
|
||
"""
|
||
conn = None
|
||
cur = None
|
||
claim_token = f"{node_code}-{int(threading.current_thread().ident or 0)}-{int(__import__('time').time())}"
|
||
try:
|
||
conn, cur = self.connect()
|
||
if not conn or not cur:
|
||
logger.error("领取检测任务失败: 无法获取数据库连接")
|
||
return []
|
||
normalized_job_id = int(job_id) if job_id not in (None, "", 0, "0") else None
|
||
normalized_limit = max(1, int(limit or 1))
|
||
normalized_lease_seconds = max(60, int(lease_seconds or 3600))
|
||
step_claim_quota = _resolve_step_claim_quota(normalized_limit)
|
||
rows = []
|
||
|
||
def _claim_rows(batch_limit: int, step_condition_sql: str = "", step_condition_params: tuple = ()) -> list:
|
||
safe_batch_limit = max(1, int(batch_limit or 0))
|
||
if safe_batch_limit <= 0:
|
||
return []
|
||
extra_condition = f"\n AND {step_condition_sql}" if step_condition_sql else ""
|
||
cur.execute(
|
||
f"""
|
||
WITH picked AS (
|
||
SELECT id
|
||
FROM detect_job_items
|
||
WHERE status IN ('pending', 'failed')
|
||
AND COALESCE(step_code, '') <> ''
|
||
AND (%s IS NULL OR job_id = %s)
|
||
AND (lease_expires_at IS NULL OR lease_expires_at < CURRENT_TIMESTAMP){extra_condition}
|
||
ORDER BY create_time ASC, id ASC
|
||
FOR UPDATE SKIP LOCKED
|
||
LIMIT %s
|
||
),
|
||
updated AS (
|
||
UPDATE detect_job_items AS item
|
||
SET status = 'claimed',
|
||
claimed_by = %s,
|
||
claim_token = %s,
|
||
lease_expires_at = CURRENT_TIMESTAMP + (%s || ' seconds')::interval,
|
||
attempt_count = item.attempt_count + 1,
|
||
started_at = COALESCE(item.started_at, CURRENT_TIMESTAMP),
|
||
updated_at = CURRENT_TIMESTAMP
|
||
FROM picked
|
||
WHERE item.id = picked.id
|
||
RETURNING item.id, item.job_id, item.domain_id, item.claim_token, item.step_code, item.step_payload_json
|
||
)
|
||
SELECT
|
||
updated.id,
|
||
updated.job_id,
|
||
updated.domain_id,
|
||
updated.claim_token,
|
||
updated.step_code,
|
||
detect_jobs.task_mode,
|
||
detect_jobs.job_code,
|
||
detect_jobs.step_code,
|
||
updated.step_payload_json,
|
||
domains.domain,
|
||
domains.source_type,
|
||
domains.register_status,
|
||
domains.detect_status,
|
||
domains.use_status,
|
||
domains.expire_date,
|
||
domains.jucha_status,
|
||
domains.juziseo_status
|
||
FROM updated
|
||
JOIN detect_jobs ON detect_jobs.id = updated.job_id
|
||
JOIN domains ON domains.id = updated.domain_id
|
||
ORDER BY updated.id ASC
|
||
""",
|
||
(
|
||
normalized_job_id,
|
||
normalized_job_id,
|
||
*step_condition_params,
|
||
safe_batch_limit,
|
||
node_code,
|
||
claim_token,
|
||
normalized_lease_seconds,
|
||
),
|
||
)
|
||
return cur.fetchall() or []
|
||
|
||
remaining = normalized_limit
|
||
for step_code in _ordered_step_claim_codes():
|
||
if remaining <= 0:
|
||
break
|
||
claimed_rows = _claim_rows(
|
||
min(remaining, step_claim_quota),
|
||
"COALESCE(step_code, '') = %s",
|
||
(str(step_code or "").strip(),),
|
||
)
|
||
if claimed_rows:
|
||
rows.extend(claimed_rows)
|
||
remaining -= len(claimed_rows)
|
||
|
||
if remaining > 0:
|
||
rows.extend(_claim_rows(remaining))
|
||
|
||
touched_job_ids = sorted({row[1] for row in rows})
|
||
for job_id in touched_job_ids:
|
||
cur.execute(
|
||
"""
|
||
UPDATE detect_jobs
|
||
SET status = CASE WHEN status = 'pending' THEN 'running' ELSE status END,
|
||
started_at = COALESCE(started_at, CURRENT_TIMESTAMP)
|
||
WHERE id = %s
|
||
""",
|
||
(job_id,),
|
||
)
|
||
conn.commit()
|
||
return [
|
||
{
|
||
"job_item_id": row[0],
|
||
"job_id": row[1],
|
||
"id": row[2],
|
||
"claim_token": row[3],
|
||
"item_step_code": row[4],
|
||
"task_mode": row[5],
|
||
"job_code": row[6],
|
||
"step_code": row[4] or row[7],
|
||
"step_payload": row[8],
|
||
"domain": row[9],
|
||
"source_type": row[10],
|
||
"register_status": row[11],
|
||
"detect_status": row[12],
|
||
"use_status": row[13],
|
||
"expire_date": row[14],
|
||
"jucha_status": row[15],
|
||
"juziseo_status": row[16],
|
||
}
|
||
for row in rows
|
||
]
|
||
except Exception as e:
|
||
try:
|
||
if conn:
|
||
conn.rollback()
|
||
except Exception:
|
||
pass
|
||
logger.error(f"领取检测任务失败: {e}")
|
||
return []
|
||
finally:
|
||
self.close(conn, cur)
|
||
|
||
def has_dispatchable_detect_job_items(self):
|
||
"""
|
||
是否仍存在可继续派发的标准步骤任务项。
|
||
|
||
这里显式排除空 step_code 的旧兼容链路项，避免 worker 在 single_step
|
||
任务已经大量积压时继续回退到旧 domains 链路补位。
|
||
"""
|
||
row = self.fetch_one(
|
||
"""
|
||
SELECT EXISTS (
|
||
SELECT 1
|
||
FROM detect_job_items
|
||
WHERE status IN ('pending', 'failed')
|
||
AND COALESCE(step_code, '') <> ''
|
||
AND (lease_expires_at IS NULL OR lease_expires_at < CURRENT_TIMESTAMP)
|
||
LIMIT 1
|
||
) AS has_items
|
||
"""
|
||
)
|
||
if isinstance(row, dict):
|
||
return bool(row.get("has_items"))
|
||
if isinstance(row, (list, tuple)) and row:
|
||
return bool(row[0])
|
||
return False
|
||
|
||
def recycle_expired_detect_job_items(self):
|
||
"""
|
||
回收租约过期但未完成的任务项，重新放回 pending。
|
||
"""
|
||
conn = None
|
||
cur = None
|
||
try:
|
||
conn, cur = self.connect()
|
||
if not conn or not cur:
|
||
return 0
|
||
cur.execute(
|
||
"""
|
||
WITH recycled AS (
|
||
UPDATE detect_job_items
|
||
SET status = 'pending',
|
||
claimed_by = '',
|
||
claim_token = '',
|
||
lease_expires_at = NULL,
|
||
updated_at = CURRENT_TIMESTAMP,
|
||
last_error = CASE
|
||
WHEN status = 'running' THEN 'lease expired while running'
|
||
WHEN status = 'claimed' THEN 'lease expired before running'
|
||
ELSE last_error
|
||
END
|
||
WHERE status IN ('claimed', 'running')
|
||
AND lease_expires_at IS NOT NULL
|
||
AND lease_expires_at < CURRENT_TIMESTAMP
|
||
RETURNING job_id
|
||
)
|
||
SELECT count(*), array_remove(array_agg(DISTINCT job_id), NULL)
|
||
FROM recycled
|
||
"""
|
||
)
|
||
row = cur.fetchone()
|
||
recycled_count = int((row or [0])[0] or 0)
|
||
touched_job_ids = (row or [0, []])[1] or []
|
||
for job_id in touched_job_ids:
|
||
self._refresh_detect_job_status_with_cursor(cur, job_id)
|
||
conn.commit()
|
||
return recycled_count
|
||
except Exception as e:
|
||
try:
|
||
if conn:
|
||
conn.rollback()
|
||
except Exception:
|
||
pass
|
||
logger.error(f"回收过期检测任务失败: {e}")
|
||
return 0
|
||
finally:
|
||
self.close(conn, cur)
|
||
|
||
def release_detect_job_items_for_node(self, node_code):
|
||
"""
|
||
释放指定节点遗留的 claimed/running 任务项，供节点重启后快速恢复。
|
||
"""
|
||
conn = None
|
||
cur = None
|
||
try:
|
||
conn, cur = self.connect()
|
||
if not conn or not cur:
|
||
return 0
|
||
cur.execute(
|
||
"""
|
||
WITH released AS (
|
||
UPDATE detect_job_items
|
||
SET status = 'pending',
|
||
claimed_by = '',
|
||
claim_token = '',
|
||
lease_expires_at = NULL,
|
||
updated_at = CURRENT_TIMESTAMP,
|
||
last_error = CASE
|
||
WHEN status = 'running' THEN 'released after worker restart'
|
||
WHEN status = 'claimed' THEN 'released before execution after worker restart'
|
||
ELSE last_error
|
||
END
|
||
WHERE claimed_by = %s
|
||
AND status IN ('claimed', 'running')
|
||
RETURNING job_id
|
||
)
|
||
SELECT count(*), array_remove(array_agg(DISTINCT job_id), NULL)
|
||
FROM released
|
||
""",
|
||
(node_code,),
|
||
)
|
||
row = cur.fetchone()
|
||
released_count = int((row or [0])[0] or 0)
|
||
touched_job_ids = (row or [0, []])[1] or []
|
||
for job_id in touched_job_ids:
|
||
self._refresh_detect_job_status_with_cursor(cur, job_id)
|
||
conn.commit()
|
||
return released_count
|
||
except Exception as e:
|
||
try:
|
||
if conn:
|
||
conn.rollback()
|
||
except Exception:
|
||
pass
|
||
logger.error(f"释放节点遗留任务失败: {e}")
|
||
return 0
|
||
finally:
|
||
self.close(conn, cur)
|
||
|
||
def get_active_detect_job(self):
|
||
"""
|
||
获取当前活动中的检测任务摘要，供 Worker 重启后自动回挂。
|
||
"""
|
||
return self.fetch_one(
|
||
"""
|
||
SELECT
|
||
job.id,
|
||
job.job_code,
|
||
job.task_mode,
|
||
job.status,
|
||
count(item.*) FILTER (WHERE item.status = 'pending') AS items_pending,
|
||
count(item.*) FILTER (WHERE item.status = 'claimed') AS items_claimed,
|
||
count(item.*) FILTER (WHERE item.status = 'running') AS items_running,
|
||
count(item.*) FILTER (WHERE item.status IN ('completed', 'blacklisted')) AS items_completed,
|
||
count(item.*) FILTER (WHERE item.status = 'failed') AS items_failed,
|
||
max(item.updated_at) AS latest_item_updated_at,
|
||
max(item.create_time) AS latest_item_created_at
|
||
FROM detect_jobs AS job
|
||
LEFT JOIN detect_job_items AS item
|
||
ON item.job_id = job.id
|
||
WHERE job.status IN ('pending', 'running')
|
||
GROUP BY job.id, job.job_code, job.task_mode, job.status
|
||
ORDER BY
|
||
CASE
|
||
WHEN count(item.*) FILTER (WHERE item.status IN ('claimed', 'running')) > 0 THEN 0
|
||
WHEN count(item.*) FILTER (WHERE item.status = 'pending') > 0 THEN 1
|
||
ELSE 2
|
||
END ASC,
|
||
CASE WHEN job.status = 'running' THEN 0 ELSE 1 END ASC,
|
||
COALESCE(max(item.updated_at), max(item.create_time), job.started_at, job.created_at) DESC,
|
||
job.id DESC
|
||
LIMIT 1
|
||
"""
|
||
)
|
||
|
||
def mark_detect_job_item_running(self, job_item_id, claim_token):
|
||
return self.execute(
|
||
"""
|
||
UPDATE detect_job_items
|
||
SET status = 'running',
|
||
updated_at = CURRENT_TIMESTAMP
|
||
WHERE id = %s AND claim_token = %s AND status = 'claimed'
|
||
""",
|
||
(job_item_id, claim_token),
|
||
)
|
||
|
||
def mark_detect_job_items_running_batch(self, items):
|
||
normalized_items = []
|
||
for item in list(items or []):
|
||
if not isinstance(item, (list, tuple)) or len(item) < 2:
|
||
continue
|
||
try:
|
||
job_item_id = int(item[0] or 0)
|
||
except Exception:
|
||
continue
|
||
claim_token = str(item[1] or "").strip()
|
||
if job_item_id <= 0 or not claim_token:
|
||
continue
|
||
normalized_items.append((job_item_id, claim_token))
|
||
if not normalized_items:
|
||
return 0
|
||
|
||
conn = None
|
||
cur = None
|
||
try:
|
||
conn, cur = self.connect()
|
||
if not conn or not cur:
|
||
return -1
|
||
values_sql = ",".join(
|
||
cur.mogrify("(%s, %s)", item).decode("utf-8")
|
||
for item in normalized_items
|
||
)
|
||
cur.execute(
|
||
f"""
|
||
UPDATE detect_job_items AS item
|
||
SET status = 'running',
|
||
updated_at = CURRENT_TIMESTAMP
|
||
FROM (VALUES {values_sql}) AS batch(id, claim_token)
|
||
WHERE item.id = batch.id
|
||
AND item.claim_token = batch.claim_token
|
||
AND item.status = 'claimed'
|
||
"""
|
||
)
|
||
updated_count = int(cur.rowcount or 0)
|
||
conn.commit()
|
||
return updated_count
|
||
except Exception as e:
|
||
try:
|
||
if conn:
|
||
conn.rollback()
|
||
except Exception:
|
||
pass
|
||
logger.error(f"批量标记检测任务项为运行中失败: {e}")
|
||
return -1
|
||
finally:
|
||
self.close(conn, cur)
|
||
|
||
def renew_detect_job_item_lease(self, job_item_id, claim_token, lease_seconds=3600):
|
||
return self.execute(
|
||
"""
|
||
UPDATE detect_job_items
|
||
SET lease_expires_at = CURRENT_TIMESTAMP + (%s || ' seconds')::interval,
|
||
updated_at = CURRENT_TIMESTAMP
|
||
WHERE id = %s AND claim_token = %s AND status IN ('claimed', 'running')
|
||
""",
|
||
(max(60, int(lease_seconds or 3600)), job_item_id, claim_token),
|
||
)
|
||
|
||
def append_detect_run_event(self, job_id, job_item_id, node_code, event_type, message, level='info', payload=None):
|
||
return self.execute(
|
||
"""
|
||
INSERT INTO detect_run_events (job_id, job_item_id, node_code, event_type, level, message, payload_json)
|
||
VALUES (%s, %s, %s, %s, %s, %s, %s)
|
||
""",
|
||
(job_id, job_item_id, node_code, event_type, level, message, Json(payload or {})),
|
||
)
|
||
|
||
def complete_detect_job_item(self, job_item_id, claim_token, final_status='completed', result_payload=None, result_version='v1', refresh_job_status=True):
|
||
conn = None
|
||
cur = None
|
||
try:
|
||
conn, cur = self.connect()
|
||
if not conn or not cur:
|
||
return False
|
||
payload_json = json.dumps(result_payload, ensure_ascii=False) if result_payload is not None else None
|
||
cur.execute(
|
||
"""
|
||
UPDATE detect_job_items
|
||
SET status = %s,
|
||
result_payload_json = COALESCE(%s::jsonb, result_payload_json),
|
||
result_version = CASE
|
||
WHEN %s <> '' THEN %s
|
||
ELSE result_version
|
||
END,
|
||
finished_at = CURRENT_TIMESTAMP,
|
||
updated_at = CURRENT_TIMESTAMP,
|
||
lease_expires_at = NULL
|
||
WHERE id = %s AND claim_token = %s
|
||
RETURNING job_id
|
||
""",
|
||
(
|
||
final_status,
|
||
payload_json,
|
||
str(result_version or '').strip(),
|
||
str(result_version or '').strip(),
|
||
job_item_id,
|
||
claim_token,
|
||
),
|
||
)
|
||
row = cur.fetchone()
|
||
if row and refresh_job_status:
|
||
self._refresh_detect_job_status_with_cursor(cur, row[0])
|
||
conn.commit()
|
||
return bool(row)
|
||
except Exception as e:
|
||
try:
|
||
if conn:
|
||
conn.rollback()
|
||
except Exception:
|
||
pass
|
||
logger.error(f"完成检测任务项失败: {e}")
|
||
return False
|
||
finally:
|
||
self.close(conn, cur)
|
||
|
||
def fail_detect_job_item(self, job_item_id, claim_token, reason='', result_payload=None, result_version='v1', refresh_job_status=True):
|
||
conn = None
|
||
cur = None
|
||
try:
|
||
conn, cur = self.connect()
|
||
if not conn or not cur:
|
||
return False
|
||
payload_json = json.dumps(result_payload, ensure_ascii=False) if result_payload is not None else None
|
||
cur.execute(
|
||
"""
|
||
UPDATE detect_job_items
|
||
SET status = 'failed',
|
||
last_error = %s,
|
||
result_payload_json = COALESCE(%s::jsonb, result_payload_json),
|
||
result_version = CASE
|
||
WHEN %s <> '' THEN %s
|
||
ELSE result_version
|
||
END,
|
||
finished_at = CURRENT_TIMESTAMP,
|
||
updated_at = CURRENT_TIMESTAMP,
|
||
lease_expires_at = NULL
|
||
WHERE id = %s AND claim_token = %s
|
||
RETURNING job_id
|
||
""",
|
||
(
|
||
str(reason or '')[:1000],
|
||
payload_json,
|
||
str(result_version or '').strip(),
|
||
str(result_version or '').strip(),
|
||
job_item_id,
|
||
claim_token,
|
||
),
|
||
)
|
||
row = cur.fetchone()
|
||
if row and refresh_job_status:
|
||
self._refresh_detect_job_status_with_cursor(cur, row[0])
|
||
conn.commit()
|
||
return bool(row)
|
||
except Exception as e:
|
||
try:
|
||
if conn:
|
||
conn.rollback()
|
||
except Exception:
|
||
pass
|
||
logger.error(f"标记检测任务项失败: {e}")
|
||
return False
|
||
finally:
|
||
self.close(conn, cur)
|
||
|
||
def finalize_detect_job_items_batch(self, items):
|
||
normalized_items = []
|
||
event_rows = []
|
||
for item in list(items or []):
|
||
if not isinstance(item, dict):
|
||
continue
|
||
try:
|
||
job_item_id = int(item.get("job_item_id") or 0)
|
||
except Exception:
|
||
continue
|
||
claim_token = str(item.get("claim_token") or "").strip()
|
||
final_status = str(item.get("final_status") or "").strip() or "failed"
|
||
message = str(item.get("message") or "").strip()
|
||
result_payload = item.get("result_payload")
|
||
result_version = str(item.get("result_version") or "v1").strip()
|
||
if job_item_id <= 0 or not claim_token:
|
||
continue
|
||
payload_json = ""
|
||
if result_payload is not None:
|
||
try:
|
||
payload_json = json.dumps(result_payload, ensure_ascii=False)
|
||
except Exception:
|
||
payload_json = ""
|
||
normalized_items.append(
|
||
(
|
||
job_item_id,
|
||
claim_token,
|
||
final_status,
|
||
message[:1000],
|
||
payload_json,
|
||
result_version,
|
||
)
|
||
)
|
||
try:
|
||
event_job_id = int(item.get("job_id") or 0)
|
||
except Exception:
|
||
event_job_id = 0
|
||
event_type = str(item.get("event_type") or "").strip()
|
||
event_message = str(item.get("event_message") or "").strip()
|
||
if event_job_id > 0 and event_type and event_message:
|
||
event_payload_json = "{}"
|
||
try:
|
||
event_payload_json = json.dumps(item.get("event_payload") or {}, ensure_ascii=False)
|
||
except Exception:
|
||
event_payload_json = "{}"
|
||
event_rows.append(
|
||
(
|
||
job_item_id,
|
||
event_job_id,
|
||
str(item.get("node_code") or "").strip(),
|
||
event_type,
|
||
str(item.get("event_level") or "info").strip() or "info",
|
||
event_message[:1000],
|
||
event_payload_json,
|
||
)
|
||
)
|
||
if not normalized_items:
|
||
return 0
|
||
|
||
conn = None
|
||
cur = None
|
||
try:
|
||
conn, cur = self.connect()
|
||
if not conn or not cur:
|
||
return -1
|
||
values_sql = ",".join(
|
||
cur.mogrify("(%s, %s, %s, %s, %s, %s)", item).decode("utf-8")
|
||
for item in normalized_items
|
||
)
|
||
cur.execute(
|
||
f"""
|
||
UPDATE detect_job_items AS item
|
||
SET status = batch.final_status,
|
||
last_error = CASE
|
||
WHEN batch.final_status = 'failed' AND batch.message <> '' THEN batch.message
|
||
ELSE item.last_error
|
||
END,
|
||
result_payload_json = COALESCE(NULLIF(batch.payload_json, '')::jsonb, item.result_payload_json),
|
||
result_version = CASE
|
||
WHEN batch.result_version <> '' THEN batch.result_version
|
||
ELSE item.result_version
|
||
END,
|
||
finished_at = CURRENT_TIMESTAMP,
|
||
updated_at = CURRENT_TIMESTAMP,
|
||
lease_expires_at = NULL
|
||
FROM (
|
||
VALUES {values_sql}
|
||
) AS batch(id, claim_token, final_status, message, payload_json, result_version)
|
||
WHERE item.id = batch.id
|
||
AND item.claim_token = batch.claim_token
|
||
AND item.status IN ('claimed', 'running')
|
||
RETURNING item.id
|
||
"""
|
||
)
|
||
updated_item_ids = {int(row[0]) for row in (cur.fetchall() or [])}
|
||
updated_count = len(updated_item_ids)
|
||
filtered_event_rows = [row for row in event_rows if int(row[0]) in updated_item_ids]
|
||
if filtered_event_rows:
|
||
values_sql = ",".join(
|
||
cur.mogrify("(%s, %s, %s, %s, %s, %s, %s::jsonb)", row).decode("utf-8")
|
||
for row in filtered_event_rows
|
||
)
|
||
cur.execute(
|
||
f"""
|
||
INSERT INTO detect_run_events (
|
||
job_item_id,
|
||
job_id,
|
||
node_code,
|
||
event_type,
|
||
level,
|
||
message,
|
||
payload_json
|
||
)
|
||
VALUES {values_sql}
|
||
"""
|
||
)
|
||
conn.commit()
|
||
return updated_count
|
||
except Exception as e:
|
||
try:
|
||
if conn:
|
||
conn.rollback()
|
||
except Exception:
|
||
pass
|
||
logger.error(f"批量完成检测任务项失败: {e}")
|
||
return -1
|
||
finally:
|
||
self.close(conn, cur)
|
||
|
||
def _refresh_detect_job_status_with_cursor(self, cur, job_id):
|
||
cur.execute(
|
||
"""
|
||
SELECT
|
||
job.task_mode,
|
||
count(*) FILTER (WHERE detect_job_items.status = 'pending') AS pending_count,
|
||
count(*) FILTER (WHERE detect_job_items.status IN ('claimed', 'running')) AS dispatch_active_count,
|
||
count(*) FILTER (WHERE detect_job_items.status = 'failed') AS failed_count,
|
||
count(*) FILTER (WHERE detect_job_items.status IN ('completed', 'blacklisted')) AS done_count,
|
||
count(*) FILTER (
|
||
WHERE detect_job_items.status IN ('completed', 'blacklisted', 'failed')
|
||
AND COALESCE(detect_job_items.step_code, '') <> ''
|
||
AND COALESCE(result_payload_json->>'controller_processed', 'false') <> 'true'
|
||
) AS unprocessed_terminal_count
|
||
FROM detect_job_items
|
||
JOIN detect_jobs AS job ON job.id = detect_job_items.job_id
|
||
WHERE job_id = %s
|
||
GROUP BY job.task_mode
|
||
""",
|
||
(job_id,),
|
||
)
|
||
row = cur.fetchone()
|
||
task_mode = str((row or [""])[0] or "").strip()
|
||
pending_count = int((row or ["", 0])[1] or 0)
|
||
dispatch_active_count = int((row or ["", 0, 0])[2] or 0)
|
||
failed_count = int((row or ["", 0, 0, 0])[3] or 0)
|
||
done_count = int((row or ["", 0, 0, 0, 0])[4] or 0)
|
||
unprocessed_terminal_count = int((row or ["", 0, 0, 0, 0, 0])[5] or 0)
|
||
if dispatch_active_count > 0 or (task_mode == 'domain_pipeline' and unprocessed_terminal_count > 0):
|
||
cur.execute(
|
||
"UPDATE detect_jobs SET status = 'running', started_at = COALESCE(started_at, CURRENT_TIMESTAMP) WHERE id = %s",
|
||
(job_id,),
|
||
)
|
||
return
|
||
if pending_count > 0:
|
||
cur.execute(
|
||
"""
|
||
UPDATE detect_jobs
|
||
SET status = 'pending',
|
||
finished_at = NULL
|
||
WHERE id = %s
|
||
""",
|
||
(job_id,),
|
||
)
|
||
return
|
||
final_status = 'completed'
|
||
if failed_count > 0 and done_count > 0:
|
||
final_status = 'partial_failed'
|
||
elif failed_count > 0:
|
||
final_status = 'failed'
|
||
cur.execute(
|
||
"""
|
||
UPDATE detect_jobs
|
||
SET status = %s,
|
||
finished_at = CURRENT_TIMESTAMP,
|
||
started_at = COALESCE(started_at, CURRENT_TIMESTAMP)
|
||
WHERE id = %s
|
||
""",
|
||
(final_status, job_id),
|
||
)
|
||
|
||
def refresh_detect_job_status(self, job_id):
|
||
conn = None
|
||
cur = None
|
||
try:
|
||
conn, cur = self.connect()
|
||
if not conn or not cur:
|
||
return False
|
||
self._refresh_detect_job_status_with_cursor(cur, job_id)
|
||
conn.commit()
|
||
return True
|
||
except Exception as e:
|
||
try:
|
||
if conn:
|
||
conn.rollback()
|
||
except Exception:
|
||
pass
|
||
logger.error(f"刷新检测任务状态失败: {e}")
|
||
return False
|
||
finally:
|
||
self.close(conn, cur)
|
||
|
||
def fetch_one(self, sql, params=None):
|
||
"""
|
||
获取单条数据
|
||
|
||
:param sql: SQL语句
|
||
:param params: 参数
|
||
:return: dict - 数据
|
||
"""
|
||
import threading
|
||
thread_id = threading.current_thread().ident
|
||
|
||
for attempt in range(2):
|
||
conn = None
|
||
cur = None
|
||
try:
|
||
conn, cur = self.connect(thread_id)
|
||
if not conn or not cur:
|
||
logger.warning(f"线程 {thread_id} 数据库连接失败，返回None")
|
||
return None
|
||
|
||
cur.execute(sql, params)
|
||
row = cur.fetchone()
|
||
if row:
|
||
columns = [desc[0] for desc in cur.description]
|
||
return dict(zip(columns, row))
|
||
return None
|
||
except (psycopg2.OperationalError, psycopg2.InterfaceError) as e:
|
||
logger.warning(f"查询数据连接异常，第 {attempt + 1} 次: {sql}, 错误: {e}")
|
||
try:
|
||
if conn and not conn.closed:
|
||
conn.close()
|
||
except Exception:
|
||
pass
|
||
if attempt == 0:
|
||
continue
|
||
logger.error(f"查询数据失败: {sql}, 错误: {e}")
|
||
return None
|
||
except Exception as e:
|
||
logger.error(f"查询数据失败: {sql}, 错误: {e}")
|
||
return None
|
||
finally:
|
||
self.close(conn, cur)
|
||
|
||
def fetch_all(self, sql, params=None):
|
||
"""
|
||
获取多条数据
|
||
|
||
:param sql: SQL语句
|
||
:param params: 参数
|
||
:return: list - 数据列表
|
||
"""
|
||
import threading
|
||
thread_id = threading.current_thread().ident
|
||
|
||
for attempt in range(2):
|
||
conn = None
|
||
cur = None
|
||
try:
|
||
conn, cur = self.connect(thread_id)
|
||
if not conn or not cur:
|
||
logger.warning(f"线程 {thread_id} 数据库连接失败，返回空结果")
|
||
return []
|
||
|
||
cur.execute(sql, params)
|
||
try:
|
||
rows = cur.fetchall()
|
||
if rows and cur.description:
|
||
columns = [desc[0] for desc in cur.description]
|
||
return [dict(zip(columns, row)) for row in rows]
|
||
return []
|
||
except Exception as e:
|
||
if "no results to fetch" in str(e):
|
||
return []
|
||
raise
|
||
except (psycopg2.OperationalError, psycopg2.InterfaceError) as e:
|
||
logger.warning(f"查询数据连接异常，第 {attempt + 1} 次: {sql}, 错误: {e}")
|
||
try:
|
||
if conn and not conn.closed:
|
||
conn.close()
|
||
except Exception:
|
||
pass
|
||
if attempt == 0:
|
||
continue
|
||
logger.error(f"查询数据失败: {sql}, 错误: {e}")
|
||
return []
|
||
except Exception as e:
|
||
logger.error(f"查询数据失败: {sql}, 错误: {e}")
|
||
return []
|
||
finally:
|
||
self.close(conn, cur)
|
||
|
||
def domain_exists(self, domain):
|
||
"""
|
||
检查域名是否存在
|
||
|
||
:param domain: 域名
|
||
:return: bool - 是否存在
|
||
"""
|
||
# 尝试使用布隆过滤器
|
||
if self.use_redis and self.use_bloom_filter:
|
||
try:
|
||
if not self.redis_client.execute_command('BF.EXISTS', 'domain_bloom', domain):
|
||
# 布隆过滤器判断不存在，直接返回 False
|
||
return False
|
||
except Exception as e:
|
||
logger.warning(f"布隆过滤器查询失败: {e}")
|
||
|
||
# 尝试使用 Redis 缓存
|
||
if self.use_redis:
|
||
try:
|
||
if self.redis_client.exists(f"domain:{domain}"):
|
||
return True
|
||
except Exception as e:
|
||
logger.warning(f"Redis 查询失败: {e}")
|
||
|
||
# 缓存未命中或 Redis 不可用，查询数据库
|
||
sql = "SELECT id FROM domains WHERE domain = %s"
|
||
result = self.fetch_one(sql, (domain,))
|
||
|
||
# 将结果存入缓存和布隆过滤器
|
||
if self.use_redis and result:
|
||
try:
|
||
pipe = self.redis_client.pipeline()
|
||
pipe.set(f"domain:{domain}", 1, ex=2592000) # 1个月过期
|
||
if self.use_bloom_filter:
|
||
pipe.execute_command('BF.ADD', 'domain_bloom', domain)
|
||
pipe.execute()
|
||
except Exception as e:
|
||
logger.warning(f"Redis 存储失败: {e}")
|
||
|
||
return result is not None
|
||
|
||
def check_domains_exist(self, domains):
|
||
"""
|
||
批量检查域名是否存在
|
||
|
||
:param domains: 域名列表
|
||
:return: list - 存在的域名列表
|
||
"""
|
||
if not domains:
|
||
return []
|
||
|
||
# 尝试使用布隆过滤器快速过滤
|
||
if self.use_redis and self.use_bloom_filter:
|
||
try:
|
||
# 分批使用布隆过滤器过滤不存在的域名
|
||
possibly_exist = []
|
||
batch_size = 10000
|
||
|
||
for i in range(0, len(domains), batch_size):
|
||
batch = domains[i:i+batch_size]
|
||
# 使用管道批量执行布隆过滤器查询
|
||
pipe = self.redis_client.pipeline()
|
||
for domain in batch:
|
||
pipe.execute_command('BF.EXISTS', 'domain_bloom', domain)
|
||
results = pipe.execute()
|
||
# 处理结果
|
||
for domain, exists in zip(batch, results):
|
||
if exists:
|
||
possibly_exist.append(domain)
|
||
|
||
# 每处理一批，记录一次进度
|
||
if (i + len(batch)) % (batch_size * 10) == 0:
|
||
logger.info(f"布隆过滤器已过滤 {min(i+len(batch), len(domains))}/{len(domains)} 个域名")
|
||
|
||
# 如果布隆过滤器判断所有域名都不存在，直接返回空列表
|
||
if not possibly_exist:
|
||
logger.info(f"布隆过滤器快速过滤: {len(domains)} 个域名不存在")
|
||
return []
|
||
|
||
# 只查询可能存在的域名
|
||
domains = possibly_exist
|
||
logger.info(f"布隆过滤器过滤后，剩余 {len(domains)} 个域名需要查询数据库")
|
||
except Exception as e:
|
||
logger.warning(f"布隆过滤器批量查询失败: {e}")
|
||
|
||
import threading
|
||
thread_id = threading.current_thread().ident
|
||
conn = None
|
||
cur = None
|
||
|
||
try:
|
||
# 从连接池获取连接
|
||
conn, cur = self.connect(thread_id)
|
||
if not conn or not cur:
|
||
# 数据库连接失败，返回空列表
|
||
logger.warning(f"线程 {thread_id} 数据库连接失败，返回空列表")
|
||
return []
|
||
|
||
# 分批查询数据库
|
||
existing_domains = []
|
||
batch_size = 10000
|
||
|
||
for i in range(0, len(domains), batch_size):
|
||
batch = domains[i:i+batch_size]
|
||
|
||
# 使用IN子句批量查询
|
||
placeholders = ','.join(['%s'] * len(batch))
|
||
sql = f"SELECT domain FROM domains WHERE domain IN ({placeholders})"
|
||
|
||
cur.execute(sql, batch)
|
||
rows = cur.fetchall()
|
||
batch_existing = [row[0] for row in rows]
|
||
existing_domains.extend(batch_existing)
|
||
|
||
# 每处理一批，记录一次进度
|
||
if (i + len(batch)) % (batch_size * 10) == 0:
|
||
logger.info(f"数据库已查询 {min(i+len(batch), len(domains))}/{len(domains)} 个域名")
|
||
|
||
# 将实际存在的域名添加到缓存，分批进行
|
||
if self.use_redis and existing_domains:
|
||
try:
|
||
batch_size = 10000
|
||
for i in range(0, len(existing_domains), batch_size):
|
||
batch = existing_domains[i:i+batch_size]
|
||
pipe = self.redis_client.pipeline()
|
||
for domain in batch:
|
||
pipe.set(f"domain:{domain}", 1, ex=3600)
|
||
pipe.execute()
|
||
except Exception as e:
|
||
logger.warning(f"Redis 批量存储失败: {e}")
|
||
|
||
logger.info(f"批量检查域名完成，发现 {len(existing_domains)} 个已存在域名")
|
||
return existing_domains
|
||
except Exception as e:
|
||
logger.error(f"批量检查域名存在失败: {e}")
|
||
return []
|
||
finally:
|
||
# 将连接放回连接池
|
||
self.close(conn, cur)
|
||
|
||
def add_domain(self, domain, tld, source_type):
|
||
"""
|
||
添加域名
|
||
|
||
:param domain: 域名
|
||
:param tld: 顶级域名
|
||
:param source_type: 来源类型
|
||
:return: int - 域名ID
|
||
"""
|
||
sql = """
|
||
INSERT INTO domains (domain, tld, source_type, use_status, detect_status, register_status, backlink_count)
|
||
VALUES (%s, %s, %s, 0, 0, 0, 0)
|
||
ON CONFLICT (domain) DO NOTHING
|
||
RETURNING id
|
||
"""
|
||
import threading
|
||
thread_id = threading.current_thread().ident
|
||
conn = None
|
||
cur = None
|
||
|
||
try:
|
||
# 从连接池获取连接
|
||
conn, cur = self.connect(thread_id)
|
||
if not conn or not cur:
|
||
# 数据库连接失败，返回None
|
||
logger.warning(f"线程 {thread_id} 数据库连接失败，返回None")
|
||
return None
|
||
|
||
cur.execute(sql, (domain, tld, source_type))
|
||
result = cur.fetchone()
|
||
conn.commit()
|
||
|
||
# 如果域名已存在，返回None
|
||
if not result:
|
||
return None
|
||
|
||
domain_id = result[0]
|
||
|
||
# 将结果存入缓存和布隆过滤器
|
||
if self.use_redis:
|
||
try:
|
||
pipe = self.redis_client.pipeline()
|
||
pipe.set(f"domain:{domain}", 1, ex=2592000) # 1个月过期
|
||
if self.use_bloom_filter:
|
||
pipe.execute_command('BF.ADD', 'domain_bloom', domain)
|
||
pipe.execute()
|
||
except Exception as e:
|
||
logger.warning(f"Redis 存储失败: {e}")
|
||
|
||
return domain_id
|
||
except Exception as e:
|
||
logger.error(f"添加域名失败: {e}")
|
||
if conn:
|
||
try:
|
||
conn.rollback()
|
||
except:
|
||
pass
|
||
return None
|
||
finally:
|
||
# 将连接放回连接池
|
||
self.close(conn, cur)
|
||
|
||
def add_domains_batch(self, domains):
|
||
"""
|
||
批量添加域名
|
||
|
||
:param domains: 域名列表，每个元素为 (domain, tld, source_type)
|
||
:return: int - 添加成功的数量
|
||
"""
|
||
if not domains:
|
||
return 0
|
||
|
||
import threading
|
||
thread_id = threading.current_thread().ident
|
||
conn = None
|
||
cur = None
|
||
|
||
try:
|
||
total_added = 0
|
||
batch_size = 1000
|
||
|
||
# 从连接池获取连接
|
||
conn, cur = self.connect(thread_id)
|
||
if not conn or not cur:
|
||
# 数据库连接失败，返回0
|
||
logger.warning(f"线程 {thread_id} 数据库连接失败，返回0")
|
||
return 0
|
||
|
||
# 分批处理
|
||
for i in range(0, len(domains), batch_size):
|
||
batch = domains[i:i+batch_size]
|
||
|
||
# 使用批量插入语法
|
||
placeholders = ','.join(['(%s, %s, %s, 0, 0, 0, 0)'] * len(batch))
|
||
sql = f"""
|
||
INSERT INTO domains (domain, tld, source_type, use_status, detect_status, register_status, backlink_count)
|
||
VALUES {placeholders}
|
||
ON CONFLICT (domain) DO NOTHING
|
||
"""
|
||
|
||
# 扁平化数据
|
||
data = []
|
||
domain_names = []
|
||
for domain, tld, source_type in batch:
|
||
data.extend([domain, tld, source_type])
|
||
domain_names.append(domain)
|
||
|
||
cur.execute(sql, data)
|
||
added_count = cur.rowcount
|
||
total_added += added_count
|
||
conn.commit()
|
||
|
||
# 为新添加的域名创建检测任务
|
||
if added_count > 0:
|
||
# 获取刚添加的域名ID
|
||
placeholders = ','.join(['%s'] * len(batch))
|
||
sql = f"SELECT id, domain FROM domains WHERE domain IN ({placeholders})"
|
||
cur.execute(sql, domain_names)
|
||
rows = cur.fetchall()
|
||
domain_ids = [row[0] for row in rows]
|
||
|
||
# 将新添加的域名添加到 Redis 缓存和布隆过滤器
|
||
if self.use_redis:
|
||
try:
|
||
pipe = self.redis_client.pipeline()
|
||
for row in rows:
|
||
domain = row[1]
|
||
pipe.set(f"domain:{domain}", 1, ex=2592000) # 1个月过期
|
||
if self.use_bloom_filter:
|
||
pipe.execute_command('BF.ADD', 'domain_bloom', domain)
|
||
pipe.execute()
|
||
except Exception as e:
|
||
logger.warning(f"Redis 批量存储失败: {e}")
|
||
|
||
# 批量创建检测任务
|
||
if domain_ids:
|
||
task_placeholders = ','.join(['(%s, 1, 0, 0, 0)'] * len(domain_ids))
|
||
task_sql = f"""
|
||
INSERT INTO detect_tasks (domain_id, task_type, status, priority, retry_count)
|
||
VALUES {task_placeholders}
|
||
"""
|
||
task_data = []
|
||
for domain_id in domain_ids:
|
||
task_data.append(domain_id)
|
||
|
||
cur.execute(task_sql, task_data)
|
||
conn.commit()
|
||
|
||
# 每处理一批，记录一次进度
|
||
if (i + len(batch)) % (batch_size * 10) == 0:
|
||
logger.info(f"已添加 {i + len(batch)}/{len(domains)} 个域名")
|
||
|
||
logger.info(f"批量添加域名完成，成功添加 {total_added} 个域名")
|
||
return total_added
|
||
except Exception as e:
|
||
logger.error(f"批量添加域名失败: {e}")
|
||
if conn:
|
||
try:
|
||
conn.rollback()
|
||
except:
|
||
pass
|
||
return 0
|
||
finally:
|
||
# 将连接放回连接池
|
||
self.close(conn, cur)
|
||
|
||
def get_domain_by_id(self, domain_id):
|
||
"""
|
||
根据ID获取域名
|
||
|
||
:param domain_id: 域名ID
|
||
:return: dict - 域名信息
|
||
"""
|
||
sql = "SELECT * FROM domains WHERE id = %s"
|
||
return self.fetch_one(sql, (domain_id,))
|
||
|
||
def get_domain_by_name(self, domain):
|
||
"""
|
||
根据域名获取信息
|
||
|
||
:param domain: 域名
|
||
:return: dict - 域名信息
|
||
"""
|
||
sql = "SELECT * FROM domains WHERE domain = %s"
|
||
return self.fetch_one(sql, (domain,))
|
||
|
||
def update_domain_use_status(self, domain_id, status):
|
||
"""
|
||
更新域名使用状态
|
||
|
||
:param domain_id: 域名ID
|
||
:param status: 状态
|
||
:return: bool - 是否更新成功
|
||
"""
|
||
sql = "UPDATE domains SET use_status = %s WHERE id = %s"
|
||
return self.execute(sql, (status, domain_id))
|
||
|
||
def update_domain_detect_status(self, domain_id, status):
|
||
"""
|
||
更新域名检测状态
|
||
|
||
:param domain_id: 域名ID
|
||
:param status: 状态
|
||
:return: bool - 是否更新成功
|
||
"""
|
||
if status == DETECT_STATUS_COMPLETED:
|
||
sql = "UPDATE domains SET detect_status = %s, detect_time = CURRENT_TIMESTAMP WHERE id = %s"
|
||
else:
|
||
sql = "UPDATE domains SET detect_status = %s WHERE id = %s"
|
||
return self.execute(sql, (status, domain_id))
|
||
|
||
def update_domain_detect_status_batch(self, items):
|
||
"""
|
||
批量更新域名检测状态，减少高并发失败风暴时的写库往返。
|
||
"""
|
||
normalized_items = []
|
||
for item in items or []:
|
||
try:
|
||
domain_id, status = item
|
||
normalized_items.append((int(domain_id), int(status)))
|
||
except Exception:
|
||
continue
|
||
if not normalized_items:
|
||
return 0
|
||
|
||
values_sql = ", ".join(["(%s, %s)"] * len(normalized_items))
|
||
params = []
|
||
for domain_id, status in normalized_items:
|
||
params.extend([domain_id, status])
|
||
|
||
sql = f"""
|
||
UPDATE domains AS d
|
||
SET detect_status = v.detect_status,
|
||
detect_time = CASE
|
||
WHEN v.detect_status = %s THEN CURRENT_TIMESTAMP
|
||
ELSE d.detect_time
|
||
END
|
||
FROM (
|
||
VALUES {values_sql}
|
||
) AS v(domain_id, detect_status)
|
||
WHERE d.id = v.domain_id
|
||
"""
|
||
params = [DETECT_STATUS_COMPLETED] + params
|
||
return len(normalized_items) if self.execute(sql, tuple(params)) else -1
|
||
|
||
def recycle_running_domains(self, target_status):
|
||
"""
|
||
回收异常中断后遗留的“检测中”状态。
|
||
|
||
:param target_status: 目标状态，通常使用 DETECT_STATUS_FAILED
|
||
:return: int - 影响行数
|
||
"""
|
||
conn, cur = self.connect()
|
||
if not conn or not cur:
|
||
logger.error("回收检测中状态失败: 无法获取数据库连接")
|
||
return 0
|
||
try:
|
||
cur.execute(
|
||
"UPDATE domains SET detect_status = %s WHERE detect_status = %s",
|
||
(target_status, DETECT_STATUS_RUNNING)
|
||
)
|
||
affected = cur.rowcount or 0
|
||
conn.commit()
|
||
return affected
|
||
except Exception as e:
|
||
try:
|
||
conn.rollback()
|
||
except Exception:
|
||
pass
|
||
logger.error(f"回收检测中状态失败: {e}")
|
||
return 0
|
||
finally:
|
||
self.close(conn, cur)
|
||
|
||
def update_domain_third_party_status(self, domain_id, field_name, status):
|
||
"""
|
||
更新第三方平台检测状态。
|
||
"""
|
||
if field_name not in {'jucha_status', 'juziseo_status'}:
|
||
logger.error(f"不支持的第三方状态字段: {field_name}")
|
||
return False
|
||
sql = f"UPDATE domains SET {field_name} = %s WHERE id = %s"
|
||
return self.execute(sql, (status, domain_id))
|
||
|
||
def mark_jucha_detected(self, domain_id):
|
||
return self.update_domain_third_party_status(domain_id, 'jucha_status', THIRD_PARTY_STATUS_DONE)
|
||
|
||
def mark_juziseo_detected(self, domain_id):
|
||
return self.update_domain_third_party_status(domain_id, 'juziseo_status', THIRD_PARTY_STATUS_DONE)
|
||
|
||
def reset_optional_detection_statuses(self, domain_id, *, jucha=False, juziseo=False):
|
||
fields = []
|
||
params = []
|
||
if jucha:
|
||
fields.append("jucha_status = 0")
|
||
if juziseo:
|
||
fields.append("juziseo_status = 0")
|
||
if not fields:
|
||
return True
|
||
sql = f"UPDATE domains SET {', '.join(fields)} WHERE id = %s"
|
||
params.append(domain_id)
|
||
return self.execute(sql, tuple(params))
|
||
|
||
def update_domain_expire_date(self, domain_id, expire_date):
|
||
"""
|
||
更新域名过期时间
|
||
|
||
:param domain_id: 域名ID
|
||
:param expire_date: 过期时间
|
||
:return: bool - 是否更新成功
|
||
"""
|
||
sql = "UPDATE domains SET expire_date = %s WHERE id = %s"
|
||
return self.execute(sql, (expire_date, domain_id))
|
||
|
||
def get_domains_to_detect(self, limit=1000, detect_options=None):
|
||
"""
|
||
获取需要检测的域名
|
||
|
||
:param limit: 限制数量
|
||
:return: list - 域名列表
|
||
"""
|
||
detect_options = detect_options or {}
|
||
conditions = [
|
||
"detect_status IN (%s, %s)",
|
||
"(use_status = 0 AND detect_status = %s AND register_status = %s AND expire_date < CURRENT_DATE)",
|
||
]
|
||
params = [
|
||
DETECT_STATUS_PENDING,
|
||
DETECT_STATUS_FAILED,
|
||
DETECT_STATUS_COMPLETED,
|
||
REGISTER_STATUS_REGISTERED,
|
||
]
|
||
|
||
if detect_options.get('detect_jucha'):
|
||
conditions.append("(detect_status <> %s AND jucha_status = 0)")
|
||
params.append(DETECT_STATUS_BLACKLISTED)
|
||
|
||
if detect_options.get('detect_juziseo'):
|
||
conditions.append("(detect_status <> %s AND juziseo_status = 0)")
|
||
params.append(DETECT_STATUS_BLACKLISTED)
|
||
|
||
sql = f"""
|
||
SELECT id, domain, source_type, register_status, detect_status, use_status, expire_date, jucha_status, juziseo_status
|
||
FROM domains
|
||
WHERE {" OR ".join(conditions)}
|
||
ORDER BY id ASC
|
||
LIMIT %s
|
||
"""
|
||
params.append(limit)
|
||
return self.fetch_all(sql, tuple(params))
|
||
|
||
def get_all_sensitive_words(self):
|
||
"""
|
||
获取所有敏感词
|
||
|
||
:return: list - 敏感词列表
|
||
"""
|
||
try:
|
||
sql = "SELECT word FROM sensitive_words"
|
||
results = self.fetch_all(sql)
|
||
sensitive_words = []
|
||
for row in results:
|
||
if isinstance(row, dict) and 'word' in row:
|
||
sensitive_words.append(row['word'])
|
||
return sensitive_words
|
||
except Exception as e:
|
||
logger.error(f"获取敏感词失败: {e}")
|
||
return []
|
||
|
||
def add_to_blacklist(self, domain, reason):
|
||
"""
|
||
将域名加入黑名单
|
||
|
||
:param domain: 域名
|
||
:param reason: 加入黑名单的原因
|
||
:return: bool - 是否操作成功
|
||
"""
|
||
sql = """
|
||
INSERT INTO blacklist (domain, reason, created_at)
|
||
VALUES (%s, %s, NOW())
|
||
ON CONFLICT (domain) DO UPDATE
|
||
SET reason = %s, updated_at = NOW()
|
||
"""
|
||
return self.execute(sql, (domain, reason, reason))
|
||
|
||
def mark_domain_blacklisted(self, domain_id, domain, reason):
|
||
"""
|
||
在一个事务中同时更新域名黑名单状态和黑名单表，减少热路径往返。
|
||
"""
|
||
conn = None
|
||
cur = None
|
||
try:
|
||
conn, cur = self.connect()
|
||
if not conn or not cur:
|
||
logger.error("标记域名黑名单失败: 无法获取数据库连接")
|
||
return False
|
||
cur.execute(
|
||
"UPDATE domains SET detect_status = %s WHERE id = %s",
|
||
(DETECT_STATUS_BLACKLISTED, domain_id),
|
||
)
|
||
cur.execute(
|
||
"""
|
||
INSERT INTO blacklist (domain, reason, created_at)
|
||
VALUES (%s, %s, NOW())
|
||
ON CONFLICT (domain) DO UPDATE
|
||
SET reason = %s, updated_at = NOW()
|
||
""",
|
||
(domain, reason, reason),
|
||
)
|
||
conn.commit()
|
||
return True
|
||
except Exception as e:
|
||
try:
|
||
if conn:
|
||
conn.rollback()
|
||
except Exception:
|
||
pass
|
||
logger.error(f"标记域名黑名单失败: {e}")
|
||
return False
|
||
finally:
|
||
self.close(conn, cur)
|
||
|
||
def update_domain_register_status(self, domain_id, status):
|
||
"""
|
||
更新域名注册状态
|
||
|
||
:param domain_id: 域名ID
|
||
:param status: 状态
|
||
:return: bool - 是否更新成功
|
||
"""
|
||
sql = "UPDATE domains SET register_status = %s WHERE id = %s"
|
||
return self.execute(sql, (status, domain_id))
|
||
|
||
def update_domain_register_result(self, domain_id, status, expire_date=None):
|
||
"""
|
||
一次性更新注册状态及过期时间，减少热路径往返次数。
|
||
|
||
:param domain_id: 域名ID
|
||
:param status: 注册状态
|
||
:param expire_date: 过期时间；为空时保持原值
|
||
:return: bool - 是否更新成功
|
||
"""
|
||
normalized_expire_date = str(expire_date).strip() if expire_date not in (None, "") else None
|
||
sql = """
|
||
UPDATE domains
|
||
SET register_status = %s,
|
||
expire_date = COALESCE(%s, expire_date)
|
||
WHERE id = %s
|
||
"""
|
||
return self.execute(sql, (status, normalized_expire_date, domain_id))
|
||
|
||
def update_domain_beian_info(self, domain_id, company_type, website_url, has_beian, beian_year):
|
||
"""
|
||
更新域名备案信息
|
||
|
||
:param domain_id: 域名ID
|
||
:param company_type: 单位性质
|
||
:param website_url: 网站首页网址
|
||
:param has_beian: 是否备案
|
||
:param beian_year: 备案年份
|
||
:return: bool - 是否更新成功
|
||
"""
|
||
sql = "UPDATE domains SET company_type = %s, website_url = %s, has_beian = %s, beian_year = %s WHERE id = %s"
|
||
return self.execute(sql, (company_type, website_url, has_beian, beian_year, domain_id))
|
||
|
||
def update_domain_beian_info_and_mark_jucha_detected(
|
||
self,
|
||
domain_id,
|
||
company_type=None,
|
||
website_url=None,
|
||
has_beian=None,
|
||
beian_year=None,
|
||
):
|
||
"""
|
||
一次性更新备案信息并标记聚查已完成，减少热路径写库次数。
|
||
"""
|
||
sql = """
|
||
UPDATE domains
|
||
SET company_type = COALESCE(%s, company_type),
|
||
website_url = COALESCE(%s, website_url),
|
||
has_beian = COALESCE(%s, has_beian),
|
||
beian_year = COALESCE(%s, beian_year),
|
||
jucha_status = %s
|
||
WHERE id = %s
|
||
"""
|
||
return self.execute(
|
||
sql,
|
||
(company_type, website_url, has_beian, beian_year, THIRD_PARTY_STATUS_DONE, domain_id),
|
||
)
|
||
|
||
def update_domain_review_status(self, domain_id, review_status):
|
||
"""
|
||
更新域名复核状态
|
||
|
||
:param domain_id: 域名ID
|
||
:param review_status: 复核状态
|
||
:return: bool - 是否更新成功
|
||
"""
|
||
sql = "UPDATE domains SET review_status = %s WHERE id = %s"
|
||
return self.execute(sql, (review_status, domain_id))
|
||
|
||
def update_domain_review_status_batch(self, items):
|
||
"""
|
||
批量更新域名复核状态。
|
||
"""
|
||
normalized_items = []
|
||
for item in items or []:
|
||
try:
|
||
domain_id, review_status = item
|
||
normalized_items.append((int(domain_id), int(review_status)))
|
||
except Exception:
|
||
continue
|
||
if not normalized_items:
|
||
return 0
|
||
|
||
values_sql = ", ".join(["(%s, %s)"] * len(normalized_items))
|
||
params = []
|
||
for domain_id, review_status in normalized_items:
|
||
params.extend([domain_id, review_status])
|
||
|
||
sql = f"""
|
||
UPDATE domains AS d
|
||
SET review_status = v.review_status
|
||
FROM (
|
||
VALUES {values_sql}
|
||
) AS v(domain_id, review_status)
|
||
WHERE d.id = v.domain_id
|
||
"""
|
||
return len(normalized_items) if self.execute(sql, tuple(params)) else -1
|
||
|
||
def update_domain_snapshot_years(self, domain_id, years):
|
||
"""
|
||
更新域名快照年份
|
||
|
||
:param domain_id: 域名ID
|
||
:param years: 年份字符串
|
||
:return: bool - 是否更新成功
|
||
"""
|
||
sql = "UPDATE domains SET snapshot_years = %s WHERE id = %s"
|
||
return self.execute(sql, (years, domain_id))
|
||
|
||
def update_domain_wayback_summary(self, domain_id, years=None, backlink_count=None):
|
||
"""
|
||
一次性更新时光机摘要字段，减少 domains 表写入次数。
|
||
"""
|
||
sql = """
|
||
UPDATE domains
|
||
SET snapshot_years = COALESCE(%s, snapshot_years),
|
||
backlink_count = COALESCE(%s, backlink_count)
|
||
WHERE id = %s
|
||
"""
|
||
return self.execute(sql, (years, backlink_count, domain_id))
|
||
|
||
def complete_domain_detection(self, domain_id, *, register_status, use_status, expire_date):
|
||
"""
|
||
一次性完成域名完成态、待复核态、过期时间归零逻辑。
|
||
|
||
:param domain_id: 域名ID
|
||
:param register_status: 当前注册状态
|
||
:param use_status: 当前使用状态
|
||
:param expire_date: 当前过期时间
|
||
:return: bool - 是否更新成功
|
||
"""
|
||
has_expire_date = bool(expire_date)
|
||
sql = """
|
||
UPDATE domains
|
||
SET detect_status = %s,
|
||
detect_time = CURRENT_TIMESTAMP,
|
||
expire_date = CASE
|
||
WHEN %s = %s AND %s = 0 AND %s THEN NULL
|
||
ELSE expire_date
|
||
END,
|
||
review_status = CASE
|
||
WHEN %s = %s THEN %s
|
||
ELSE review_status
|
||
END
|
||
WHERE id = %s
|
||
"""
|
||
return self.execute(
|
||
sql,
|
||
(
|
||
DETECT_STATUS_COMPLETED,
|
||
register_status,
|
||
REGISTER_STATUS_AVAILABLE,
|
||
use_status,
|
||
has_expire_date,
|
||
register_status,
|
||
REGISTER_STATUS_AVAILABLE,
|
||
REVIEW_STATUS_PENDING,
|
||
domain_id,
|
||
),
|
||
)
|
||
|
||
def complete_domain_detection_batch(self, items):
|
||
"""
|
||
批量更新域名完成态，减少 completed 尾部的单条写库往返。
|
||
"""
|
||
normalized_items = []
|
||
for item in items or []:
|
||
try:
|
||
domain_id, register_status, use_status, has_expire_date = item
|
||
normalized_items.append(
|
||
(
|
||
int(domain_id),
|
||
int(register_status),
|
||
int(use_status),
|
||
bool(has_expire_date),
|
||
)
|
||
)
|
||
except Exception:
|
||
continue
|
||
if not normalized_items:
|
||
return 0
|
||
|
||
values_sql = ", ".join(["(%s, %s, %s, %s)"] * len(normalized_items))
|
||
params = []
|
||
for domain_id, register_status, use_status, has_expire_date in normalized_items:
|
||
params.extend([domain_id, register_status, use_status, has_expire_date])
|
||
|
||
sql = f"""
|
||
UPDATE domains AS d
|
||
SET detect_status = %s,
|
||
detect_time = CURRENT_TIMESTAMP,
|
||
expire_date = CASE
|
||
WHEN v.register_status = %s AND v.use_status = 0 AND v.has_expire_date THEN NULL
|
||
ELSE d.expire_date
|
||
END,
|
||
review_status = CASE
|
||
WHEN v.register_status = %s THEN %s
|
||
ELSE d.review_status
|
||
END
|
||
FROM (
|
||
VALUES {values_sql}
|
||
) AS v(domain_id, register_status, use_status, has_expire_date)
|
||
WHERE d.id = v.domain_id
|
||
"""
|
||
params = [
|
||
DETECT_STATUS_COMPLETED,
|
||
REGISTER_STATUS_AVAILABLE,
|
||
REGISTER_STATUS_AVAILABLE,
|
||
REVIEW_STATUS_PENDING,
|
||
] + params
|
||
return len(normalized_items) if self.execute(sql, tuple(params)) else -1
|
||
|
||
def create_detect_task(self, domain_id, task_type, priority=0):
|
||
"""
|
||
创建检测任务
|
||
|
||
:param domain_id: 域名ID
|
||
:param task_type: 任务类型
|
||
:param priority: 优先级
|
||
:return: int - 任务ID
|
||
"""
|
||
sql = """
|
||
INSERT INTO detect_tasks (domain_id, task_type, status, priority, retry_count)
|
||
VALUES (%s, %s, 0, %s, 0)
|
||
RETURNING id
|
||
"""
|
||
import threading
|
||
thread_id = threading.current_thread().ident
|
||
conn = None
|
||
cur = None
|
||
try:
|
||
conn, cur = self.connect(thread_id)
|
||
if not conn or not cur:
|
||
return None
|
||
cur.execute(sql, (domain_id, task_type, priority))
|
||
row = cur.fetchone()
|
||
conn.commit()
|
||
return row[0] if row else None
|
||
except Exception as e:
|
||
logger.error(f"创建检测任务失败: {e}")
|
||
if conn:
|
||
conn.rollback()
|
||
return None
|
||
finally:
|
||
self.close(conn, cur)
|
||
|
||
def get_pending_task(self):
|
||
"""
|
||
获取待执行的任务
|
||
|
||
:return: dict - 任务信息
|
||
"""
|
||
sql = """
|
||
SELECT * FROM detect_tasks
|
||
WHERE status = 0
|
||
ORDER BY priority DESC, create_time ASC
|
||
LIMIT 1
|
||
"""
|
||
return self.fetch_one(sql)
|
||
|
||
def get_task_by_id(self, task_id):
|
||
"""
|
||
根据ID获取任务
|
||
|
||
:param task_id: 任务ID
|
||
:return: dict - 任务信息
|
||
"""
|
||
sql = "SELECT * FROM detect_tasks WHERE id = %s"
|
||
return self.fetch_one(sql, (task_id,))
|
||
|
||
def update_task_status(self, task_id, status):
|
||
"""
|
||
更新任务状态
|
||
|
||
:param task_id: 任务ID
|
||
:param status: 状态
|
||
:return: bool - 是否更新成功
|
||
"""
|
||
sql = "UPDATE detect_tasks SET status = %s WHERE id = %s"
|
||
return self.execute(sql, (status, task_id))
|
||
|
||
def update_task_retry_count(self, task_id, retry_count):
|
||
"""
|
||
更新任务重试次数
|
||
|
||
:param task_id: 任务ID
|
||
:param retry_count: 重试次数
|
||
:return: bool - 是否更新成功
|
||
"""
|
||
sql = "UPDATE detect_tasks SET retry_count = %s WHERE id = %s"
|
||
return self.execute(sql, (retry_count, task_id))
|
||
|
||
def get_failed_tasks(self):
|
||
"""
|
||
获取失败的任务
|
||
|
||
:return: list - 任务列表
|
||
"""
|
||
sql = "SELECT * FROM detect_tasks WHERE status = 3"
|
||
return self.fetch_all(sql)
|
||
|
||
def clear_completed_tasks(self, days):
|
||
"""
|
||
清理已完成的任务
|
||
|
||
:param days: 保留天数
|
||
:return: int - 清理的任务数量
|
||
"""
|
||
sql = "DELETE FROM detect_tasks WHERE status = 2 AND create_time < NOW() - (%s * INTERVAL '1 day')"
|
||
import threading
|
||
thread_id = threading.current_thread().ident
|
||
conn = None
|
||
cur = None
|
||
try:
|
||
conn, cur = self.connect(thread_id)
|
||
if not conn or not cur:
|
||
return 0
|
||
cur.execute(sql, (days,))
|
||
count = cur.rowcount
|
||
conn.commit()
|
||
return count
|
||
except Exception as e:
|
||
logger.error(f"清理已完成任务失败: {e}")
|
||
if conn:
|
||
conn.rollback()
|
||
return 0
|
||
finally:
|
||
self.close(conn, cur)
|
||
|
||
def add_to_blacklist(self, domain, reason):
|
||
"""
|
||
添加到黑名单
|
||
|
||
:param domain: 域名
|
||
:param reason: 原因
|
||
:return: bool - 是否添加成功
|
||
"""
|
||
sql = """
|
||
INSERT INTO domain_blacklist (domain, reason)
|
||
VALUES (%s, %s)
|
||
ON CONFLICT (domain) DO NOTHING
|
||
"""
|
||
return self.execute(sql, (domain, reason))
|
||
|
||
def is_blacklisted(self, domain):
|
||
"""
|
||
检查域名是否在黑名单中
|
||
|
||
:param domain: 域名
|
||
:return: bool - 是否在黑名单中
|
||
"""
|
||
sql = "SELECT id FROM domain_blacklist WHERE domain = %s"
|
||
result = self.fetch_one(sql, (domain,))
|
||
return result is not None
|
||
|
||
def add_detection_result(self, domain_id, baidu_history, baidu_site, qihu360_site, google_site, chinaz_info, aizhan_info, juziseo_info, jucha_info):
|
||
"""
|
||
添加检测结果
|
||
|
||
:param domain_id: 域名ID
|
||
:param baidu_history: 百度历史
|
||
:param baidu_site: 百度site
|
||
:param qihu360_site: 360 site
|
||
:param google_site: Google site
|
||
:param chinaz_info: 站长之家信息
|
||
:param aizhan_info: 爱站网信息
|
||
:param juziseo_info: 桔子SEO信息
|
||
:param jucha_info: 聚查信息
|
||
:return: bool - 是否添加成功
|
||
"""
|
||
import threading
|
||
thread_id = threading.current_thread().ident
|
||
conn = None
|
||
cur = None
|
||
try:
|
||
conn, cur = self.connect(thread_id)
|
||
if not conn or not cur:
|
||
return False
|
||
payload = (
|
||
json.dumps(baidu_history) if baidu_history is not None else None,
|
||
json.dumps(baidu_site) if baidu_site is not None else None,
|
||
json.dumps(qihu360_site) if qihu360_site is not None else None,
|
||
json.dumps(google_site) if google_site is not None else None,
|
||
json.dumps(chinaz_info) if chinaz_info is not None else None,
|
||
json.dumps(aizhan_info) if aizhan_info is not None else None,
|
||
json.dumps(juziseo_info) if juziseo_info is not None else None,
|
||
json.dumps(jucha_info) if jucha_info is not None else None,
|
||
)
|
||
cur.execute("SELECT id FROM domain_detections WHERE domain_id = %s ORDER BY id ASC LIMIT 1", (domain_id,))
|
||
exists = cur.fetchone()
|
||
if exists:
|
||
sql = """
|
||
UPDATE domain_detections
|
||
SET baidu_history = %s, baidu_site = %s, qihu360_site = %s, google_site = %s,
|
||
chinaz_info = %s, aizhan_info = %s, juziseo_info = %s, jucha_info = %s,
|
||
update_time = CURRENT_TIMESTAMP
|
||
WHERE domain_id = %s
|
||
"""
|
||
cur.execute(sql, payload + (domain_id,))
|
||
else:
|
||
sql = """
|
||
INSERT INTO domain_detections (
|
||
domain_id, baidu_history, baidu_site, qihu360_site, google_site,
|
||
chinaz_info, aizhan_info, juziseo_info, jucha_info
|
||
)
|
||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
|
||
"""
|
||
cur.execute(sql, (domain_id,) + payload)
|
||
conn.commit()
|
||
return True
|
||
except Exception as e:
|
||
logger.error(f"添加检测结果失败: {e}")
|
||
if conn:
|
||
conn.rollback()
|
||
return False
|
||
finally:
|
||
self.close(conn, cur)
|
||
|
||
def get_domains_by_conditions(self, conditions, page=1, page_size=1000):
|
||
"""
|
||
根据条件获取域名
|
||
|
||
:param conditions: 条件
|
||
:param page: 页码，从1开始
|
||
:param page_size: 每页数量
|
||
:return: list - 域名列表
|
||
"""
|
||
# 构建SQL语句
|
||
sql = "SELECT * FROM domains WHERE 1=1"
|
||
params = []
|
||
|
||
if 'register_status' in conditions:
|
||
sql += " AND register_status = %s"
|
||
params.append(conditions['register_status'])
|
||
|
||
if 'use_status' in conditions:
|
||
sql += " AND use_status = %s"
|
||
params.append(conditions['use_status'])
|
||
|
||
if 'detect_status' in conditions:
|
||
sql += " AND detect_status = %s"
|
||
params.append(conditions['detect_status'])
|
||
|
||
if 'review_status' in conditions:
|
||
sql += " AND review_status = %s"
|
||
params.append(conditions['review_status'])
|
||
|
||
if 'has_beian' in conditions:
|
||
sql += " AND has_beian = %s"
|
||
params.append(conditions['has_beian'])
|
||
|
||
if conditions.get('company_type'):
|
||
sql += " AND company_type LIKE %s"
|
||
params.append(f"%{conditions['company_type']}%")
|
||
|
||
if conditions.get('beian_year'):
|
||
sql += " AND beian_year = %s"
|
||
params.append(conditions['beian_year'])
|
||
|
||
if conditions.get('snapshot_year'):
|
||
sql += " AND snapshot_years LIKE %s"
|
||
params.append(f"%{conditions['snapshot_year']}%")
|
||
|
||
# 域名搜索
|
||
if conditions.get('domain'):
|
||
sql += " AND domain LIKE %s"
|
||
params.append(f"%{conditions['domain']}%")
|
||
|
||
# 计算偏移量
|
||
offset = (page - 1) * page_size
|
||
sql += " LIMIT %s OFFSET %s"
|
||
params.extend([page_size, offset])
|
||
|
||
return self.fetch_all(sql, params)
|
||
|
||
def get_domains_count(self, conditions):
|
||
"""
|
||
获取符合条件的域名总数
|
||
|
||
:param conditions: 条件
|
||
:return: int - 域名总数
|
||
"""
|
||
# 构建SQL语句
|
||
sql = """
|
||
SELECT COUNT(DISTINCT d.id)
|
||
FROM domains d
|
||
LEFT JOIN domain_detections dd ON d.id = dd.domain_id
|
||
WHERE 1=1
|
||
"""
|
||
params = []
|
||
|
||
# 只添加非空条件
|
||
if 'register_status' in conditions and conditions['register_status'] is not None:
|
||
sql += " AND d.register_status = %s"
|
||
params.append(conditions['register_status'])
|
||
|
||
if 'use_status' in conditions and conditions['use_status'] is not None:
|
||
sql += " AND d.use_status = %s"
|
||
params.append(conditions['use_status'])
|
||
|
||
if 'detect_status' in conditions and conditions['detect_status'] is not None:
|
||
sql += " AND d.detect_status = %s"
|
||
params.append(conditions['detect_status'])
|
||
|
||
if 'review_status' in conditions and conditions['review_status'] is not None:
|
||
sql += " AND d.review_status = %s"
|
||
params.append(conditions['review_status'])
|
||
|
||
if 'has_beian' in conditions and conditions['has_beian'] is not None:
|
||
sql += " AND d.has_beian = %s"
|
||
params.append(conditions['has_beian'])
|
||
|
||
# 其他条件保持不变
|
||
if conditions.get('beian_year'):
|
||
sql += " AND d.beian_year = %s"
|
||
params.append(conditions['beian_year'])
|
||
|
||
if conditions.get('snapshot_year'):
|
||
sql += " AND d.snapshot_years LIKE %s"
|
||
params.append(f"%{conditions['snapshot_year']}%")
|
||
|
||
# 域名搜索
|
||
if conditions.get('search_keyword'):
|
||
sql += " AND d.domain LIKE %s"
|
||
params.append(f"%{conditions['search_keyword']}%")
|
||
|
||
# 首页网址搜索
|
||
if conditions.get('website_url'):
|
||
sql += " AND d.website_url LIKE %s"
|
||
params.append(f"%{conditions['website_url']}%")
|
||
|
||
if conditions.get('backlink_gt_10') is True:
|
||
sql += " AND COALESCE(dd.backlink_count_gt_10, FALSE) = TRUE"
|
||
|
||
# 打印查询信息
|
||
logger.info(f"执行计数SQL: {sql}")
|
||
logger.info(f"计数参数: {params}")
|
||
|
||
result = self.fetch_one(sql, params)
|
||
if result:
|
||
count = result.get('count', 0)
|
||
logger.info(f"符合条件的域名总数: {count}")
|
||
return count
|
||
return 0
|
||
|
||
def get_domains_with_details(self, conditions, page=1, page_size=1000):
|
||
"""
|
||
获取域名及其详细信息
|
||
|
||
:param conditions: 条件
|
||
:param page: 页码，从1开始
|
||
:param page_size: 每页数量
|
||
:return: list - 域名列表
|
||
"""
|
||
# 构建SQL语句
|
||
sql = """
|
||
SELECT DISTINCT d.*, dd.baidu_site, dd.google_site, dd.qihu360_site, dd.baidu_history,
|
||
dd.chinaz_info, dd.aizhan_info, dd.juziseo_info, dd.jucha_info,
|
||
dd.is_chinese_title, dd.same_url, dd.backlink_count_gt_10
|
||
FROM domains d
|
||
LEFT JOIN domain_detections dd ON d.id = dd.domain_id
|
||
WHERE 1=1
|
||
"""
|
||
params = []
|
||
|
||
# 只添加非空条件
|
||
if 'register_status' in conditions and conditions['register_status'] is not None:
|
||
sql += " AND d.register_status = %s"
|
||
params.append(conditions['register_status'])
|
||
|
||
if 'use_status' in conditions and conditions['use_status'] is not None:
|
||
sql += " AND d.use_status = %s"
|
||
params.append(conditions['use_status'])
|
||
|
||
if 'detect_status' in conditions and conditions['detect_status'] is not None:
|
||
sql += " AND d.detect_status = %s"
|
||
params.append(conditions['detect_status'])
|
||
|
||
if 'review_status' in conditions and conditions['review_status'] is not None:
|
||
sql += " AND d.review_status = %s"
|
||
params.append(conditions['review_status'])
|
||
|
||
if 'has_beian' in conditions and conditions['has_beian'] is not None:
|
||
sql += " AND d.has_beian = %s"
|
||
params.append(conditions['has_beian'])
|
||
|
||
# 其他条件保持不变
|
||
if conditions.get('beian_year'):
|
||
sql += " AND d.beian_year = %s"
|
||
params.append(conditions['beian_year'])
|
||
|
||
if conditions.get('snapshot_year'):
|
||
sql += " AND d.snapshot_years LIKE %s"
|
||
params.append(f"%{conditions['snapshot_year']}%")
|
||
|
||
# 域名搜索
|
||
if conditions.get('search_keyword'):
|
||
sql += " AND d.domain LIKE %s"
|
||
params.append(f"%{conditions['search_keyword']}%")
|
||
|
||
# 首页网址搜索
|
||
if conditions.get('website_url'):
|
||
sql += " AND d.website_url LIKE %s"
|
||
params.append(f"%{conditions['website_url']}%")
|
||
|
||
if conditions.get('backlink_gt_10') is True:
|
||
sql += " AND COALESCE(dd.backlink_count_gt_10, FALSE) = TRUE"
|
||
|
||
sql += " ORDER BY d.id ASC"
|
||
|
||
# 计算偏移量
|
||
offset = (page - 1) * page_size
|
||
sql += " LIMIT %s OFFSET %s"
|
||
params.extend([page_size, offset])
|
||
|
||
# 打印查询信息
|
||
logger.info(f"执行查询SQL: {sql}")
|
||
logger.info(f"查询参数: {params}")
|
||
|
||
result = self.fetch_all(sql, params)
|
||
logger.info(f"查询结果数量: {len(result)}")
|
||
|
||
# 如果没有结果，尝试执行一个简单的查询来检查数据库是否有数据
|
||
if not result:
|
||
simple_sql = "SELECT COUNT(*) FROM domains"
|
||
count_result = self.fetch_one(simple_sql)
|
||
if count_result:
|
||
logger.info(f"数据库中总域名数量: {count_result.get('count', 0)}")
|
||
else:
|
||
logger.warning("无法获取数据库中域名数量")
|
||
|
||
return result
|
||
|
||
def get_domain_statistics(self):
|
||
"""
|
||
获取域名统计信息
|
||
|
||
:return: dict - 统计信息
|
||
"""
|
||
sql = """
|
||
SELECT
|
||
COUNT(*) AS total,
|
||
SUM(CASE WHEN register_status = %s THEN 1 ELSE 0 END) AS available,
|
||
SUM(CASE WHEN register_status = %s THEN 1 ELSE 0 END) AS registered,
|
||
SUM(CASE WHEN detect_status = %s THEN 1 ELSE 0 END) AS blacklisted
|
||
FROM domains
|
||
"""
|
||
result = self.fetch_one(sql, (REGISTER_STATUS_AVAILABLE, REGISTER_STATUS_REGISTERED, DETECT_STATUS_BLACKLISTED))
|
||
if result:
|
||
return {
|
||
'total': result.get('total', 0),
|
||
'available': result.get('available', 0),
|
||
'registered': result.get('registered', 0),
|
||
'blacklisted': result.get('blacklisted', 0)
|
||
}
|
||
return {}
|
||
|
||
def get_task_statistics(self):
|
||
"""
|
||
获取任务统计信息
|
||
|
||
:return: dict - 统计信息
|
||
"""
|
||
sql = """
|
||
SELECT
|
||
COUNT(*) AS total,
|
||
SUM(CASE WHEN status = 0 THEN 1 ELSE 0 END) AS pending,
|
||
SUM(CASE WHEN status = 1 THEN 1 ELSE 0 END) AS running,
|
||
SUM(CASE WHEN status = 2 THEN 1 ELSE 0 END) AS completed,
|
||
SUM(CASE WHEN status = 3 THEN 1 ELSE 0 END) AS failed
|
||
FROM detect_tasks
|
||
"""
|
||
result = self.fetch_one(sql)
|
||
if result:
|
||
return {
|
||
'total': result.get('total', 0),
|
||
'pending': result.get('pending', 0),
|
||
'running': result.get('running', 0),
|
||
'completed': result.get('completed', 0),
|
||
'failed': result.get('failed', 0)
|
||
}
|
||
return {}
|
||
|
||
def update_domain_status(self, domain_id, status_type, status_value):
|
||
"""
|
||
更新域名状态
|
||
|
||
:param domain_id: 域名ID
|
||
:param status_type: 状态类型
|
||
:param status_value: 状态值
|
||
:return: bool - 是否更新成功
|
||
"""
|
||
if status_type == 'use_status':
|
||
return self.update_domain_use_status(domain_id, status_value)
|
||
elif status_type == 'detect_status':
|
||
return self.update_domain_detect_status(domain_id, status_value)
|
||
elif status_type == 'register_status':
|
||
return self.update_domain_register_status(domain_id, status_value)
|
||
else:
|
||
logger.error(f"未知的状态类型: {status_type}")
|
||
return False
|
||
|
||
def is_ykj_domain(self, domain_id):
|
||
"""
|
||
检查域名是否为一口价域名
|
||
|
||
:param domain_id: 域名ID
|
||
:return: bool - 是否为一口价域名
|
||
"""
|
||
import threading
|
||
thread_id = threading.current_thread().ident
|
||
conn = None
|
||
cur = None
|
||
|
||
try:
|
||
conn, cur = self.connect(thread_id)
|
||
if not conn or not cur:
|
||
logger.warning(f"线程 {thread_id} 数据库连接失败，返回默认值")
|
||
return False
|
||
|
||
sql = "SELECT source_type FROM domains WHERE id = %s"
|
||
cur.execute(sql, (domain_id,))
|
||
try:
|
||
result = cur.fetchone()
|
||
if result:
|
||
# 1 表示聚名一口价
|
||
return result[0] == 1
|
||
return False
|
||
except Exception as e:
|
||
# 处理查询结果为空的情况
|
||
if "no results to fetch" in str(e):
|
||
return False
|
||
raise
|
||
except Exception as e:
|
||
logger.error(f"检查一口价域名出错: {e}")
|
||
return False
|
||
finally:
|
||
# 将连接放回连接池
|
||
self.close(conn, cur)
|