admin/getDomain
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
main
getDomain/domainCheck/app/utils/database.py

4310 lines
163 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
# -*- coding: UTF-8 -*-
'''
@Project :domainScanDemo
@File :database.py
@IDE :PyCharm
@Author :梦伴
@Date :2026/4/9 0:07
@explain : 数据库操作类
'''
import hashlib
import json
import os
import socket
import threading
import time
import uuid
from datetime import datetime, timedelta
from pathlib import Path
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.redis_client import get_redis_client
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))
_DETECT_JOB_ITEM_RECYCLE_BATCH_SIZE = 2000
_DETECT_JOB_ITEM_RECYCLE_MAX_BATCHES = 6
_DETECT_JOB_ITEM_RELEASE_BATCH_SIZE = 2000
_DETECT_JOB_ITEM_RELEASE_MAX_BATCHES = 6
_RUNTIME_INDEX_BUILD_LOCK_SCOPE = "detect-job-items-runtime-indexes"
_RUNTIME_REQUIRED_INDEX_DDL = {
"idx_detect_job_items_job_domain_step": """
CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS idx_detect_job_items_job_domain_step
ON detect_job_items(job_id, domain_id, step_code)
""",
"idx_detect_job_items_claim_step_ready": """
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_detect_job_items_claim_step_ready
ON detect_job_items(status, step_code, lease_expires_at, create_time, id)
WHERE step_code <> '' AND status IN ('pending', 'failed')
""",
"idx_detect_job_items_claim_job_step_ready": """
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_detect_job_items_claim_job_step_ready
ON detect_job_items(job_id, status, step_code, lease_expires_at, create_time, id)
WHERE step_code <> '' AND status IN ('pending', 'failed')
""",
"idx_detect_job_items_stalled_job_activity": """
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_detect_job_items_stalled_job_activity
ON detect_job_items(
job_id,
status,
(COALESCE(updated_at, started_at, create_time)),
id
)
WHERE status IN ('claimed', 'running')
""",
"idx_detect_job_items_release_node_job": """
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_detect_job_items_release_node_job
ON detect_job_items(claimed_by, job_id, status, id)
WHERE claimed_by <> ''
AND status IN ('claimed', 'running')
""",
"idx_detect_sync_records_scope_created": """
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_detect_sync_records_scope_created
ON detect_sync_records(sync_type, source_region, target_region, created_at DESC, id DESC)
""",
"idx_detect_sync_records_source_record_created": """
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_detect_sync_records_source_record_created
ON detect_sync_records(
sync_type,
source_region,
target_region,
((payload_json->>'source_record_id')),
created_at DESC,
id DESC
)
""",
"idx_detect_sync_records_source_record_hash_created": """
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_detect_sync_records_source_record_hash_created
ON detect_sync_records(
sync_type,
source_region,
target_region,
((payload_json->>'source_record_id')),
((payload_json->>'projection_hash')),
created_at DESC,
id DESC
)
""",
"idx_detect_sync_records_runtime_push_lookup": """
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_detect_sync_records_runtime_push_lookup
ON detect_sync_records(
source_region,
target_region,
((payload_json->>'sync_type')),
((payload_json->>'source_record_id')),
created_at DESC,
id DESC
)
WHERE sync_type = 'runtime_push'
""",
}
_RUNTIME_REQUIRED_INDEX_TABLES = {
"idx_detect_job_items_job_domain_step": "detect_job_items",
"idx_detect_job_items_claim_step_ready": "detect_job_items",
"idx_detect_job_items_claim_job_step_ready": "detect_job_items",
"idx_detect_job_items_stalled_job_activity": "detect_job_items",
"idx_detect_job_items_release_node_job": "detect_job_items",
"idx_detect_sync_records_scope_created": "detect_sync_records",
"idx_detect_sync_records_source_record_created": "detect_sync_records",
"idx_detect_sync_records_source_record_hash_created": "detect_sync_records",
"idx_detect_sync_records_runtime_push_lookup": "detect_sync_records",
}
_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 _select_preferred_claim_job_ids(job_rows, *, limit=8, recent_hours=24, now=None):
safe_limit = max(1, int(limit or 1))
safe_recent_hours = max(1, int(recent_hours or 1))
normalized_rows = []
for row in list(job_rows or []):
if not isinstance(row, (list, tuple)) or len(row) < 3:
continue
try:
job_id = int(row[0] or 0)
except Exception:
continue
if job_id <= 0:
continue
status = str(row[1] or "").strip().lower()
activity_at = row[2]
if activity_at is None:
continue
reference_now = now
if reference_now is None:
reference_now = (
datetime.now(activity_at.tzinfo)
if getattr(activity_at, "tzinfo", None) is not None
else datetime.now()
)
if status != "running" and activity_at < reference_now - timedelta(hours=safe_recent_hours):
continue
normalized_rows.append((job_id, status, activity_at))
normalized_rows.sort(
key=lambda item: (
0 if item[1] == "running" else 1,
-item[2].timestamp(),
-item[0],
),
)
preferred_job_ids = []
seen_job_ids = set()
for job_id, _status, _activity_at in normalized_rows:
if job_id in seen_job_ids:
continue
preferred_job_ids.append(job_id)
seen_job_ids.add(job_id)
if len(preferred_job_ids) >= safe_limit:
break
return preferred_job_ids
def _build_pg_advisory_lock_key(scope: str) -> int:
normalized_scope = str(scope or "").strip() or "domaincheck-default"
digest = hashlib.sha1(normalized_scope.encode("utf-8")).digest()
raw_value = int.from_bytes(digest[:8], "big", signed=False)
return raw_value - (1 << 64) if raw_value >= (1 << 63) else raw_value
def _build_claim_token(node_code, thread_id=None):
normalized_node_code = str(node_code or "").strip() or "unknown"
normalized_thread_id = int(
thread_id if thread_id is not None else (threading.current_thread().ident or 0)
)
node_fragment = normalized_node_code.replace(" ", "_")[:24]
thread_fragment = format(normalized_thread_id & 0xFFFFFFFF, "x")
entropy = (
f"{normalized_node_code}|{normalized_thread_id}|"
f"{time.time_ns()}|{uuid.uuid4().hex}"
)
digest = hashlib.sha1(entropy.encode("utf-8")).hexdigest()[:24]
return f"{node_fragment}-{thread_fragment}-{digest}"[:64]
def _read_local_json_config(filename, default):
for root in _local_config_search_roots():
candidate = root / str(filename or "").strip()
try:
if candidate.exists():
with candidate.open("r", encoding="utf-8") as handle:
return json.load(handle)
except Exception:
continue
return default
def _local_config_search_roots():
seen = set()
roots = []
def append_root(raw_path):
text = str(raw_path or "").strip()
if not text:
return
path = Path(text).expanduser().resolve()
key = str(path)
if key in seen:
return
seen.add(key)
roots.append(path)
explicit_root = str(os.getenv("DOMAINCHECK_CONFIG_ROOT", "") or "").strip()
append_root(explicit_root)
append_root(Path.cwd())
module_root = Path(__file__).resolve().parents[2]
append_root(module_root)
for base_root in list(roots):
normalized = str(base_root)
if f"{os.sep}releases{os.sep}" in normalized:
install_root = normalized.split(f"{os.sep}releases{os.sep}", 1)[0]
append_root(Path(install_root) / "current" / "domainCheck")
append_root(Path(install_root) / "domainCheck")
elif f"{os.sep}current{os.sep}" in normalized:
install_root = normalized.split(f"{os.sep}current{os.sep}", 1)[0]
append_root(Path(install_root) / "current" / "domainCheck")
append_root(Path(install_root) / "domainCheck")
elif base_root.name == "domainCheck":
append_root(base_root.parent / "current" / "domainCheck")
append_root(base_root.parent / "domainCheck")
return roots
def _normalize_scaling_int(value, default):
try:
normalized = int(value)
except Exception:
return int(default)
return max(1, normalized)
def _resolve_scaling_override(overrides, *candidate_codes):
if not isinstance(overrides, dict):
return None
for candidate in candidate_codes:
normalized_candidate = str(candidate or "").strip()
if not normalized_candidate:
continue
if normalized_candidate not in overrides:
continue
try:
return max(1, int(overrides.get(normalized_candidate) or 0))
except Exception:
continue
return None
def _load_local_worker_scaling_hints():
node_code = str(getattr(config, "NODE_CODE", "") or "").strip()
parent_node_code = str(os.getenv("WORKER_PARENT_NODE_CODE", "") or "").strip()
thread_count_payload = _read_local_json_config("thread_count.json", {"thread_count": "1000"})
node_thread_counts = _read_local_json_config("node_thread_counts.json", {})
process_count_payload = _read_local_json_config("process_count.json", {"process_count": "1"})
node_process_counts = _read_local_json_config("node_process_counts.json", {})
default_thread_count = _normalize_scaling_int(thread_count_payload.get("thread_count", 1000), 1000)
default_process_count = _normalize_scaling_int(process_count_payload.get("process_count", 1), 1)
effective_thread_count = _resolve_scaling_override(node_thread_counts, node_code, parent_node_code)
if effective_thread_count is None:
effective_thread_count = default_thread_count
effective_process_count = _resolve_scaling_override(node_process_counts, node_code, parent_node_code)
if effective_process_count is None:
effective_process_count = default_process_count
return {
"node_code": node_code,
"parent_node_code": parent_node_code,
"thread_count": effective_thread_count,
"process_count": effective_process_count,
}
def _is_controller_scaling_hints(scaling_hints):
normalized_hints = dict(scaling_hints or {})
node_role = str(getattr(config, "NODE_ROLE", "") or "").strip().lower()
if node_role == "control":
return True
node_code = str(normalized_hints.get("node_code") or "").strip().lower()
parent_node_code = str(normalized_hints.get("parent_node_code") or "").strip().lower()
return "controller" in node_code or "controller" in parent_node_code
def _default_pool_total_budget(process_count, scaling_hints=None):
normalized = max(1, int(process_count or 1))
if _is_controller_scaling_hints(scaling_hints):
if normalized >= 96:
return 1000
if normalized >= 48:
return 720
if normalized >= 24:
return 480
if normalized >= 48:
return 240
if normalized >= 24:
return 320
return 480
def _resolve_db_pool_limits():
scaling_hints = _load_local_worker_scaling_hints()
raw_pool_size = str(os.getenv("DB_POOL_SIZE", "") or "").strip()
raw_pool_warm_size = str(os.getenv("DB_POOL_WARM_SIZE", "") or "").strip()
raw_pool_idle_keep_max = str(os.getenv("DB_POOL_IDLE_KEEP_MAX", "") or "").strip()
effective_process_count = max(1, int(scaling_hints.get("process_count", 1) or 1))
effective_thread_count = max(1, int(scaling_hints.get("thread_count", 1000) or 1000))
if raw_pool_size:
pool_size = max(1, int(raw_pool_size))
else:
default_budget = _default_pool_total_budget(
effective_process_count,
scaling_hints=scaling_hints,
)
total_budget = max(32, int(os.getenv("DOMAINCHECK_DB_POOL_TOTAL_BUDGET", str(default_budget)) or default_budget))
per_process_budget = max(4, total_budget // effective_process_count)
thread_window = max(8, min(24, max(1, effective_thread_count // 96)))
pool_size = max(4, min(24, min(per_process_budget, thread_window)))
if raw_pool_warm_size:
pool_warm_size = max(1, min(pool_size, int(raw_pool_warm_size)))
else:
pool_warm_size = max(1, min(pool_size, max(1, min(4, pool_size // 4))))
if raw_pool_idle_keep_max:
pool_idle_keep_max = max(pool_warm_size, min(pool_size, int(raw_pool_idle_keep_max)))
else:
pool_idle_keep_max = max(pool_warm_size, min(pool_size, max(2, pool_size // 2)))
return {
"pool_size": pool_size,
"pool_warm_size": pool_warm_size,
"pool_idle_keep_max": pool_idle_keep_max,
"scaling_hints": scaling_hints,
}
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))
# Large controller pools should be allowed to claim a full-step window;
# otherwise a 2k-thread worker can get pinned near ~500 inflight items when
# a single hot step dominates the queue.
if normalized_limit >= 1024:
return normalized_limit
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
pool_limits = _resolve_db_pool_limits()
# 数据库连接池
self.connection_pool = []
self.pool_size = max(1, int(pool_limits["pool_size"])) # 连接池大小
self.pool_warm_size = max(1, min(self.pool_size, int(pool_limits["pool_warm_size"])))
self.pool_idle_keep_max = max(
self.pool_warm_size,
min(self.pool_size, int(pool_limits["pool_idle_keep_max"])),
)
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 = {}
scaling_hints = dict(pool_limits.get("scaling_hints") or {})
logger.info(
"数据库连接池限额: "
f"node={scaling_hints.get('node_code') or '-'} "
f"parent={scaling_hints.get('parent_node_code') or '-'} "
f"process_count={int(scaling_hints.get('process_count', 1) or 1)} "
f"thread_count={int(scaling_hints.get('thread_count', 0) or 0)} "
f"pool_size={self.pool_size} warm={self.pool_warm_size} idle_keep={self.pool_idle_keep_max}"
)
# 初始化连接池
self._init_connection_pool()
# 初始化 Redis 客户端
try:
self.redis_client = get_redis_client(role="standard")
# 测试连接
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
self._active_detect_job_cache_lock = threading.Lock()
self._active_detect_job_cache_payload = None
self._active_detect_job_cache_fresh_until = 0.0
self._active_detect_job_cache_stale_until = 0.0
def _clone_active_detect_job_payload(self, payload):
if isinstance(payload, dict):
return dict(payload)
return payload
def _ensure_active_detect_job_cache_state(self):
if not hasattr(self, "_active_detect_job_cache_lock") or self._active_detect_job_cache_lock is None:
self._active_detect_job_cache_lock = threading.Lock()
if not hasattr(self, "_active_detect_job_cache_payload"):
self._active_detect_job_cache_payload = None
if not hasattr(self, "_active_detect_job_cache_fresh_until"):
self._active_detect_job_cache_fresh_until = 0.0
if not hasattr(self, "_active_detect_job_cache_stale_until"):
self._active_detect_job_cache_stale_until = 0.0
def _get_local_active_detect_job_cache(self, now_ts: float, *, allow_stale: bool = False):
self._ensure_active_detect_job_cache_state()
with self._active_detect_job_cache_lock:
payload = self._clone_active_detect_job_payload(self._active_detect_job_cache_payload)
fresh_until = float(getattr(self, "_active_detect_job_cache_fresh_until", 0.0) or 0.0)
stale_until = float(getattr(self, "_active_detect_job_cache_stale_until", 0.0) or 0.0)
if payload is None:
return None
if now_ts <= fresh_until:
return payload
if allow_stale and now_ts <= stale_until:
return payload
return None
def _set_local_active_detect_job_cache(self, payload, *, now_ts: float, fresh_ttl_seconds: int, stale_ttl_seconds: int):
self._ensure_active_detect_job_cache_state()
fresh_until = now_ts + max(1, int(fresh_ttl_seconds or 1))
stale_until = fresh_until + max(0, int(stale_ttl_seconds or 0))
with self._active_detect_job_cache_lock:
self._active_detect_job_cache_payload = self._clone_active_detect_job_payload(payload)
self._active_detect_job_cache_fresh_until = fresh_until
self._active_detect_job_cache_stale_until = stale_until
def _read_cached_active_detect_job_from_redis(self, cache_key: str):
if not self.redis_client:
return None
try:
cached_payload = self.redis_client.get(cache_key)
if cached_payload:
cached = json.loads(cached_payload)
if isinstance(cached, dict):
return cached
except Exception as e:
logger.debug(f"读取活动任务缓存失败: {e}")
return None
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 _prepare_pooled_connection(self, conn):
if not conn or conn.closed:
return False
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()
try:
cur.execute("SELECT 1")
cur.fetchone()
finally:
try:
cur.close()
except Exception:
pass
self._connection_last_healthcheck[id(conn)] = now
return True
except Exception:
return False
def _discard_connection(self, conn):
conn_id = id(conn) if conn is not None else 0
try:
if conn and not conn.closed:
conn.close()
except Exception:
pass
with self.pool_condition:
self._connection_last_healthcheck.pop(conn_id, None)
self.total_connections = max(0, self.total_connections - 1)
self.pool_condition.notify()
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
pooled_conn = None
with self.pool_condition:
while self.connection_pool:
pooled_conn = self.connection_pool.pop()
if pooled_conn:
break
if pooled_conn is None and self.total_connections < self.pool_size:
self.total_connections += 1
create_new = True
elif pooled_conn is None:
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 pooled_conn is not None:
if self._prepare_pooled_connection(pooled_conn):
if _POOL_VERBOSE_LOGS:
logger.debug(f"线程 {thread_id} 从连接池获取连接成功")
return pooled_conn, pooled_conn.cursor()
self._discard_connection(pooled_conn)
continue
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:
try:
if conn.get_transaction_status() != extensions.TRANSACTION_STATUS_IDLE:
conn.rollback()
except Exception:
self._discard_connection(conn)
logger.warning("连接归还前回滚失败，已关闭连接")
return
close_conn = False
conn_id = id(conn)
with self.pool_condition:
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(conn_id, None)
self.total_connections = max(0, self.total_connections - 1)
self.pool_condition.notify()
close_conn = True
if _POOL_VERBOSE_LOGS:
logger.debug(f"空闲连接超过保留阈值({self.pool_idle_keep_max})，已关闭多余连接")
if close_conn:
try:
conn.close()
except Exception:
pass
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):
if self._cluster_runtime_schema_ready():
return True
if self._cluster_runtime_schema_basics_present():
if not self._runtime_index_repair_enabled():
logger.warning("多机运行库索引存在缺口/无效，默认跳过自动DDL修复")
return False
self._ensure_cluster_runtime_indexes()
return self._cluster_runtime_schema_ready()
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);
CREATE INDEX IF NOT EXISTS idx_detect_job_items_claim_ready
ON detect_job_items(status, create_time, id)
WHERE step_code <> '';
CREATE INDEX IF NOT EXISTS idx_detect_job_items_claim_job_ready
ON detect_job_items(job_id, status, create_time, id)
WHERE step_code <> '';
CREATE INDEX IF NOT EXISTS idx_detect_job_items_claim_step_ready
ON detect_job_items(status, step_code, lease_expires_at, create_time, id)
WHERE step_code <> '' AND status IN ('pending', 'failed');
CREATE INDEX IF NOT EXISTS idx_detect_job_items_claim_job_step_ready
ON detect_job_items(job_id, status, step_code, lease_expires_at, create_time, id)
WHERE step_code <> '' AND status IN ('pending', 'failed');
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 _cluster_runtime_schema_basics_present(self):
sql = """
SELECT
EXISTS (
SELECT 1
FROM information_schema.tables
WHERE table_schema = current_schema()
AND table_name = 'detect_worker_nodes'
) AS has_detect_worker_nodes,
EXISTS (
SELECT 1
FROM information_schema.tables
WHERE table_schema = current_schema()
AND table_name = 'detect_jobs'
) AS has_detect_jobs,
EXISTS (
SELECT 1
FROM information_schema.tables
WHERE table_schema = current_schema()
AND table_name = 'detect_job_items'
) AS has_detect_job_items,
EXISTS (
SELECT 1
FROM information_schema.tables
WHERE table_schema = current_schema()
AND table_name = 'detect_run_events'
) AS has_detect_run_events,
EXISTS (
SELECT 1
FROM information_schema.tables
WHERE table_schema = current_schema()
AND table_name = 'detect_sync_records'
) AS has_detect_sync_records,
EXISTS (
SELECT 1
FROM information_schema.columns
WHERE table_schema = current_schema()
AND table_name = 'detect_jobs'
AND column_name = 'task_mode'
) AS has_detect_jobs_task_mode,
EXISTS (
SELECT 1
FROM information_schema.columns
WHERE table_schema = current_schema()
AND table_name = 'detect_jobs'
AND column_name = 'step_code'
) AS has_detect_jobs_step_code,
EXISTS (
SELECT 1
FROM information_schema.columns
WHERE table_schema = current_schema()
AND table_name = 'detect_job_items'
AND column_name = 'step_code'
) AS has_detect_job_items_step_code,
EXISTS (
SELECT 1
FROM information_schema.columns
WHERE table_schema = current_schema()
AND table_name = 'detect_job_items'
AND column_name = 'step_payload_json'
) AS has_detect_job_items_step_payload,
EXISTS (
SELECT 1
FROM information_schema.columns
WHERE table_schema = current_schema()
AND table_name = 'detect_job_items'
AND column_name = 'result_payload_json'
) AS has_detect_job_items_result_payload
"""
try:
row = self.fetch_one(sql)
except Exception as e:
logger.debug(f"检测多机运行库 schema 状态失败: {e}")
return False
if isinstance(row, dict):
return all(bool(value) for value in row.values())
if isinstance(row, (list, tuple)):
return all(bool(value) for value in row)
return False
def _cluster_runtime_missing_indexes(self):
table_names = sorted(set(_RUNTIME_REQUIRED_INDEX_TABLES.values()))
sql = """
SELECT
idx.relname AS index_name,
pg_index.indisvalid AS is_valid
FROM pg_class AS idx
JOIN pg_index ON pg_index.indexrelid = idx.oid
JOIN pg_class AS tbl ON tbl.oid = pg_index.indrelid
JOIN pg_namespace AS ns ON ns.oid = tbl.relnamespace
WHERE ns.nspname = current_schema()
AND tbl.relname = ANY(%s)
AND idx.relname = ANY(%s)
"""
try:
rows = self.fetch_all(sql, (table_names, list(_RUNTIME_REQUIRED_INDEX_DDL.keys()),))
except Exception as e:
logger.debug(f"检测多机运行库索引状态失败: {e}")
return list(_RUNTIME_REQUIRED_INDEX_DDL.keys())
states = {
index_name: False
for index_name in _RUNTIME_REQUIRED_INDEX_DDL.keys()
}
for row in list(rows or []):
index_name = str((row or {}).get("index_name") or "").strip()
if index_name in states:
states[index_name] = bool((row or {}).get("is_valid"))
return [index_name for index_name, is_valid in states.items() if not is_valid]
def _cluster_runtime_invalid_indexes(self):
table_names = sorted(set(_RUNTIME_REQUIRED_INDEX_TABLES.values()))
sql = """
SELECT
idx.relname AS index_name
FROM pg_class AS idx
JOIN pg_index ON pg_index.indexrelid = idx.oid
JOIN pg_class AS tbl ON tbl.oid = pg_index.indrelid
JOIN pg_namespace AS ns ON ns.oid = tbl.relnamespace
WHERE ns.nspname = current_schema()
AND tbl.relname = ANY(%s)
AND idx.relname = ANY(%s)
AND NOT pg_index.indisvalid
"""
try:
rows = self.fetch_all(sql, (table_names, list(_RUNTIME_REQUIRED_INDEX_DDL.keys()),))
except Exception as e:
logger.debug(f"检测多机运行库无效索引状态失败: {e}")
return []
return [
str((row or {}).get("index_name") or "").strip()
for row in list(rows or [])
if str((row or {}).get("index_name") or "").strip() in _RUNTIME_REQUIRED_INDEX_DDL
]
def _runtime_index_repair_enabled(self):
raw = str(os.getenv("DOMAINCHECK_RUNTIME_INDEX_REPAIR_ENABLED", "") or "").strip().lower()
return raw in {"1", "true", "yes", "on"}
def _ensure_cluster_runtime_indexes(self):
missing_indexes = self._cluster_runtime_missing_indexes()
if not missing_indexes:
return True
conn = None
cur = None
lock_acquired = False
try:
conn = self._create_connection()
conn.autocommit = True
cur = conn.cursor()
cur.execute("SELECT pg_try_advisory_lock(%s)", (_build_pg_advisory_lock_key(_RUNTIME_INDEX_BUILD_LOCK_SCOPE),))
lock_row = cur.fetchone()
lock_acquired = bool((lock_row or [False])[0])
if not lock_acquired:
logger.info("多机运行库索引补齐进行中，当前进程跳过重复建索引")
return False
current_missing = set(self._cluster_runtime_missing_indexes())
current_invalid = set(self._cluster_runtime_invalid_indexes())
for index_name in missing_indexes:
if index_name not in current_missing:
continue
ddl = str(_RUNTIME_REQUIRED_INDEX_DDL.get(index_name) or "").strip()
if not ddl:
continue
if index_name in current_invalid:
if not self._runtime_index_repair_enabled():
logger.warning(f"检测到无效多机运行库索引，跳过自动重建: {index_name}")
continue
logger.warning(f"检测到无效多机运行库索引，准备重建: {index_name}")
cur.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {index_name}")
logger.info(f"补齐多机运行库索引: {index_name}")
cur.execute(ddl)
return True
except Exception as e:
logger.warning(f"补齐多机运行库索引失败: {e}")
return False
finally:
if cur and lock_acquired:
try:
cur.execute("SELECT pg_advisory_unlock(%s)", (_build_pg_advisory_lock_key(_RUNTIME_INDEX_BUILD_LOCK_SCOPE),))
except Exception:
pass
try:
if cur:
cur.close()
except Exception:
pass
try:
if conn:
conn.close()
except Exception:
pass
def _cluster_runtime_schema_ready(self):
return self._cluster_runtime_schema_basics_present() and not self._cluster_runtime_missing_indexes()
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,
*,
prefer_recent_jobs=False,
preferred_recent_job_limit=8,
preferred_recent_job_window_hours=24,
):
"""
领取一批待执行的任务项。
"""
conn = None
cur = None
claim_token = _build_claim_token(node_code)
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 = []
preferred_job_ids = []
if normalized_job_id is None and bool(prefer_recent_jobs):
cur.execute(
"""
SELECT id, status, COALESCE(started_at, created_at) AS activity_at
FROM detect_jobs
WHERE status IN ('pending', 'running')
ORDER BY
COALESCE(started_at, created_at) DESC,
CASE WHEN status = 'running' THEN 0 ELSE 1 END ASC,
id DESC
LIMIT %s
""",
(max(8, int(preferred_recent_job_limit or 8) * 4),),
)
preferred_job_ids = _select_preferred_claim_job_ids(
cur.fetchall() or [],
limit=preferred_recent_job_limit,
recent_hours=preferred_recent_job_window_hours,
)
def _claim_rows(
batch_limit: int,
step_condition_sql: str = "",
step_condition_params: tuple = (),
job_scope_sql: str = "",
job_scope_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 ""
job_scope_clause = f"\n AND {job_scope_sql}" if job_scope_sql else ""
cur.execute(
f"""
WITH picked AS (
SELECT id
FROM detect_job_items
WHERE status IN ('pending', 'failed')
AND step_code <> ''
{job_scope_clause}
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
""",
(
*job_scope_params,
*step_condition_params,
safe_batch_limit,
node_code,
claim_token,
normalized_lease_seconds,
),
)
return cur.fetchall() or []
def _claim_with_job_scope(job_scope_sql: str = "", job_scope_params: tuple = ()) -> list:
scoped_rows = []
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(),),
job_scope_sql,
job_scope_params,
)
if claimed_rows:
scoped_rows.extend(claimed_rows)
remaining -= len(claimed_rows)
if remaining > 0:
scoped_rows.extend(_claim_rows(remaining, job_scope_sql=job_scope_sql, job_scope_params=job_scope_params))
return scoped_rows
if normalized_job_id is not None:
rows.extend(_claim_with_job_scope("job_id = %s", (normalized_job_id,)))
else:
if preferred_job_ids:
rows.extend(_claim_with_job_scope("job_id = ANY(%s)", (preferred_job_ids,)))
remaining = max(0, normalized_limit - len(rows))
if remaining > 0:
original_limit = normalized_limit
try:
normalized_limit = remaining
rows.extend(_claim_with_job_scope())
finally:
normalized_limit = original_limit
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 claim_restart_released_detect_job_items(
self,
node_code,
job_id,
*,
limit=1000,
lease_seconds=3600,
):
"""
优先回收节点重启后放回 pending 的当前 job 尾批任务项。
这批条目如果长期不被重新 claim，会让旧 job 维持 running，
从而拖住下一批的自然接棒。
"""
conn = None
cur = None
claim_token = _build_claim_token(node_code)
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
if not normalized_job_id:
return []
normalized_limit = max(1, int(limit or 1))
normalized_lease_seconds = max(60, int(lease_seconds or 3600))
cur.execute(
"""
WITH picked AS (
SELECT id
FROM detect_job_items
WHERE job_id = %s
AND status = 'pending'
AND step_code <> ''
AND (lease_expires_at IS NULL OR lease_expires_at < CURRENT_TIMESTAMP)
AND last_error IN (
'released after worker restart',
'released before execution after worker restart'
)
ORDER BY
CASE
WHEN last_error = 'released after worker restart' THEN 0
ELSE 1
END ASC,
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_limit,
node_code,
claim_token,
normalized_lease_seconds,
),
)
rows = cur.fetchall() or []
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
advisory_lock_key = _build_pg_advisory_lock_key("detect-job-items-recycle-expired")
try:
conn, cur = self.connect()
if not conn or not cur:
return 0
cur.execute("SELECT pg_try_advisory_lock(%s)", (advisory_lock_key,))
lock_row = cur.fetchone()
if not bool((lock_row or [False])[0]):
try:
conn.rollback()
except Exception:
pass
return 0
recycled_count = 0
touched_job_ids = set()
for _ in range(_DETECT_JOB_ITEM_RECYCLE_MAX_BATCHES):
cur.execute(
"""
WITH expired_candidates AS (
SELECT id, job_id, status
FROM detect_job_items
WHERE status IN ('claimed', 'running')
AND lease_expires_at IS NOT NULL
AND lease_expires_at < CURRENT_TIMESTAMP
ORDER BY lease_expires_at ASC, id ASC
FOR UPDATE SKIP LOCKED
LIMIT %s
),
recycled AS (
UPDATE detect_job_items AS item
SET status = 'pending',
claimed_by = '',
claim_token = '',
lease_expires_at = NULL,
updated_at = CURRENT_TIMESTAMP,
last_error = CASE
WHEN expired_candidates.status = 'running' THEN 'lease expired while running'
WHEN expired_candidates.status = 'claimed' THEN 'lease expired before running'
ELSE item.last_error
END
FROM expired_candidates
WHERE item.id = expired_candidates.id
RETURNING expired_candidates.job_id
)
SELECT count(*), array_remove(array_agg(DISTINCT job_id), NULL)
FROM recycled
""",
(_DETECT_JOB_ITEM_RECYCLE_BATCH_SIZE,),
)
row = cur.fetchone()
batch_count = int((row or [0])[0] or 0)
recycled_count += batch_count
touched_job_ids.update(item for item in ((row or [0, []])[1] or []) if item is not None)
if batch_count < _DETECT_JOB_ITEM_RECYCLE_BATCH_SIZE:
break
for job_id in sorted(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:
if cur:
try:
cur.execute("SELECT pg_advisory_unlock(%s)", (advisory_lock_key,))
if conn:
conn.commit()
except Exception:
try:
if conn:
conn.rollback()
except Exception:
pass
self.close(conn, cur)
def recycle_stalled_detect_job_items(self, job_id, *, stall_seconds=1800, batch_size=None):
"""
定向回收长时间无活动的 claimed/running 任务项。
"""
conn = None
cur = None
try:
normalized_job_id = int(job_id or 0)
except Exception:
normalized_job_id = 0
if normalized_job_id <= 0:
return 0
safe_stall_seconds = max(300, int(stall_seconds or 1800))
safe_batch_size = max(
1,
min(
int(batch_size or _DETECT_JOB_ITEM_RECYCLE_BATCH_SIZE),
_DETECT_JOB_ITEM_RECYCLE_BATCH_SIZE,
),
)
advisory_lock_key = _build_pg_advisory_lock_key(
f"detect-job-items-recycle-stalled:{normalized_job_id}"
)
try:
conn, cur = self.connect()
if not conn or not cur:
return 0
cur.execute("SELECT pg_try_advisory_lock(%s)", (advisory_lock_key,))
lock_row = cur.fetchone()
if not bool((lock_row or [False])[0]):
try:
conn.rollback()
except Exception:
pass
return 0
cur.execute(
"""
WITH stalled_candidates AS (
SELECT id, job_id, status
FROM detect_job_items
WHERE job_id = %s
AND status IN ('claimed', 'running')
AND COALESCE(updated_at, started_at, create_time)
< CURRENT_TIMESTAMP - (%s || ' seconds')::interval
ORDER BY COALESCE(updated_at, started_at, create_time) ASC, id ASC
FOR UPDATE SKIP LOCKED
LIMIT %s
),
recycled AS (
UPDATE detect_job_items AS item
SET status = 'pending',
claimed_by = '',
claim_token = '',
lease_expires_at = NULL,
updated_at = CURRENT_TIMESTAMP,
last_error = CASE
WHEN stalled_candidates.status = 'running' THEN 'stalled while running'
WHEN stalled_candidates.status = 'claimed' THEN 'stalled before running'
ELSE item.last_error
END
FROM stalled_candidates
WHERE item.id = stalled_candidates.id
RETURNING stalled_candidates.job_id
)
SELECT count(*), array_remove(array_agg(DISTINCT job_id), NULL)
FROM recycled
""",
(
normalized_job_id,
safe_stall_seconds,
safe_batch_size,
),
)
row = cur.fetchone()
recycled_count = int((row or [0])[0] or 0)
touched_job_ids = [item for item in ((row or [0, []])[1] or []) if item is not None]
for touched_job_id in sorted(set(touched_job_ids)):
self._refresh_detect_job_status_with_cursor(cur, touched_job_id)
conn.commit()
return recycled_count
except Exception as e:
try:
if conn:
conn.rollback()
except Exception:
pass
logger.error(f"回收长挂检测任务失败: job_id={normalized_job_id}, error={e}")
return 0
finally:
if cur:
try:
cur.execute("SELECT pg_advisory_unlock(%s)", (advisory_lock_key,))
if conn:
conn.commit()
except Exception:
try:
if conn:
conn.rollback()
except Exception:
pass
self.close(conn, cur)
def release_detect_job_item(self, job_item_id, claim_token, reason=''):
"""
释放单个 claimed/running 任务项，供会话切换时快速回到 pending。
"""
conn = None
cur = None
try:
conn, cur = self.connect()
if not conn or not cur:
return False
cur.execute(
"""
UPDATE detect_job_items
SET status = 'pending',
claimed_by = '',
claim_token = '',
lease_expires_at = NULL,
updated_at = CURRENT_TIMESTAMP,
last_error = CASE
WHEN %s <> '' THEN %s
ELSE last_error
END
WHERE id = %s
AND claim_token = %s
AND status IN ('claimed', 'running')
RETURNING job_id
""",
(
str(reason or '')[:1000],
str(reason or '')[:1000],
int(job_item_id or 0),
str(claim_token or '').strip(),
),
)
row = cur.fetchone()
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 release_detect_job_items_batch(self, items):
"""
批量释放 claimed/running 任务项，降低会话切换时的逐条写锁竞争。
"""
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()
reason = ""
if len(item) >= 3:
reason = str(item[2] or "").strip()[:1000]
if job_item_id <= 0 or not claim_token:
continue
normalized_items.append((job_item_id, claim_token, reason))
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)", item).decode("utf-8")
for item in normalized_items
)
cur.execute(
f"""
UPDATE detect_job_items AS item
SET status = 'pending',
claimed_by = '',
claim_token = '',
lease_expires_at = NULL,
updated_at = CURRENT_TIMESTAMP,
last_error = CASE
WHEN batch.reason <> '' THEN batch.reason
ELSE item.last_error
END
FROM (VALUES {values_sql}) AS batch(id, claim_token, reason)
WHERE item.id = batch.id
AND item.claim_token = batch.claim_token
AND item.status IN ('claimed', 'running')
RETURNING item.id
"""
)
updated_rows = cur.fetchall() or []
conn.commit()
return len(updated_rows)
except Exception as e:
try:
if conn:
conn.rollback()
except Exception:
pass
logger.error(f"批量释放检测任务项失败: {e}")
return -1
finally:
self.close(conn, cur)
def release_detect_job_items_for_node(self, node_code):
"""
释放指定节点遗留的 claimed/running 任务项，供节点重启后快速恢复。
"""
normalized_flag = str(os.getenv("DOMAINCHECK_ENABLE_NODE_ITEM_RELEASE", "1") or "1").strip().lower()
if normalized_flag in {
"0",
"false",
"no",
"off",
"disable",
"disabled",
}:
return 0
conn = None
cur = None
try:
conn, cur = self.connect()
if not conn or not cur:
return 0
released_count = 0
touched_job_ids = set()
for _ in range(_DETECT_JOB_ITEM_RELEASE_MAX_BATCHES):
cur.execute(
"""
WITH release_candidates AS (
SELECT id, job_id, status
FROM detect_job_items
WHERE claimed_by = %s
AND status IN ('claimed', 'running')
ORDER BY id ASC
FOR UPDATE SKIP LOCKED
LIMIT %s
),
released AS (
UPDATE detect_job_items AS item
SET status = 'pending',
claimed_by = '',
claim_token = '',
lease_expires_at = NULL,
updated_at = CURRENT_TIMESTAMP,
last_error = CASE
WHEN release_candidates.status = 'running' THEN 'released after worker restart'
WHEN release_candidates.status = 'claimed' THEN 'released before execution after worker restart'
ELSE item.last_error
END
FROM release_candidates
WHERE item.id = release_candidates.id
RETURNING release_candidates.job_id
)
SELECT count(*), array_remove(array_agg(DISTINCT job_id), NULL)
FROM released
""",
(node_code, _DETECT_JOB_ITEM_RELEASE_BATCH_SIZE),
)
row = cur.fetchone()
batch_count = int((row or [0])[0] or 0)
released_count += batch_count
touched_job_ids.update(item for item in ((row or [0, []])[1] or []) if item is not None)
if batch_count < _DETECT_JOB_ITEM_RELEASE_BATCH_SIZE:
break
for job_id in sorted(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 release_detect_job_items_for_node_job(self, node_code, job_id):
"""
只释放指定节点在某个 job 上遗留的 claimed/running 任务项，供尾批定向回收。
"""
normalized_flag = str(os.getenv("DOMAINCHECK_ENABLE_NODE_ITEM_RELEASE", "1") or "1").strip().lower()
if normalized_flag in {
"0",
"false",
"no",
"off",
"disable",
"disabled",
}:
return 0
normalized_node_code = str(node_code or "").strip()
try:
normalized_job_id = int(job_id or 0)
except Exception:
normalized_job_id = 0
if not normalized_node_code or normalized_job_id <= 0:
return 0
conn = None
cur = None
try:
conn, cur = self.connect()
if not conn or not cur:
return 0
released_count = 0
touched_job_ids = set()
for _ in range(_DETECT_JOB_ITEM_RELEASE_MAX_BATCHES):
cur.execute(
"""
WITH release_candidates AS (
SELECT id, job_id, status
FROM detect_job_items
WHERE claimed_by = %s
AND job_id = %s
AND status IN ('claimed', 'running')
ORDER BY id ASC
FOR UPDATE SKIP LOCKED
LIMIT %s
),
released AS (
UPDATE detect_job_items AS item
SET status = 'pending',
claimed_by = '',
claim_token = '',
lease_expires_at = NULL,
updated_at = CURRENT_TIMESTAMP,
last_error = CASE
WHEN release_candidates.status = 'running' THEN 'released after worker restart'
WHEN release_candidates.status = 'claimed' THEN 'released before execution after worker restart'
ELSE item.last_error
END
FROM release_candidates
WHERE item.id = release_candidates.id
RETURNING release_candidates.job_id
)
SELECT count(*), array_remove(array_agg(DISTINCT job_id), NULL)
FROM released
""",
(
normalized_node_code,
normalized_job_id,
_DETECT_JOB_ITEM_RELEASE_BATCH_SIZE,
),
)
row = cur.fetchone()
batch_count = int((row or [0])[0] or 0)
released_count += batch_count
touched_job_ids.update(item for item in ((row or [0, []])[1] or []) if item is not None)
if batch_count < _DETECT_JOB_ITEM_RELEASE_BATCH_SIZE:
break
for touched_job_id in sorted(touched_job_ids):
self._refresh_detect_job_status_with_cursor(cur, touched_job_id)
conn.commit()
return released_count
except Exception as e:
try:
if conn:
conn.rollback()
except Exception:
pass
logger.error(
f"释放节点定向任务失败: node={normalized_node_code}, job_id={normalized_job_id}, error={e}"
)
return 0
finally:
self.close(conn, cur)
def get_active_detect_job(self):
"""
获取当前活动中的检测任务摘要，供 Worker 重启后自动回挂。
"""
cache_key = "domaincheck:active_detect_job_summary:v1"
cache_lock_key = f"{cache_key}:refresh_lock"
now_ts = time.time()
try:
cache_ttl_seconds = max(
1,
int(os.getenv("DOMAINCHECK_ACTIVE_JOB_CACHE_TTL_SECONDS", "3") or 3),
)
except (TypeError, ValueError):
cache_ttl_seconds = 3
try:
local_cache_ttl_seconds = max(
1,
int(
os.getenv(
"DOMAINCHECK_ACTIVE_JOB_LOCAL_CACHE_TTL_SECONDS",
str(cache_ttl_seconds),
)
or cache_ttl_seconds
),
)
except (TypeError, ValueError):
local_cache_ttl_seconds = cache_ttl_seconds
try:
stale_cache_ttl_seconds = max(
0,
int(os.getenv("DOMAINCHECK_ACTIVE_JOB_CACHE_STALE_SECONDS", "10") or 10),
)
except (TypeError, ValueError):
stale_cache_ttl_seconds = 10
try:
cache_refresh_lock_seconds = max(
1,
int(os.getenv("DOMAINCHECK_ACTIVE_JOB_CACHE_LOCK_SECONDS", "2") or 2),
)
except (TypeError, ValueError):
cache_refresh_lock_seconds = 2
local_cached = self._get_local_active_detect_job_cache(now_ts)
if local_cached is not None:
return local_cached
cached = self._read_cached_active_detect_job_from_redis(cache_key)
if isinstance(cached, dict):
self._set_local_active_detect_job_cache(
cached,
now_ts=now_ts,
fresh_ttl_seconds=local_cache_ttl_seconds,
stale_ttl_seconds=stale_cache_ttl_seconds,
)
return self._clone_active_detect_job_payload(cached)
refresh_lock_token = ""
if self.redis_client:
try:
refresh_lock_token = str(uuid.uuid4())
acquired = self.redis_client.set(
cache_lock_key,
refresh_lock_token,
nx=True,
ex=cache_refresh_lock_seconds,
)
if not acquired:
stale_cached = self._get_local_active_detect_job_cache(now_ts, allow_stale=True)
if stale_cached is not None:
return stale_cached
time.sleep(0.05)
cached = self._read_cached_active_detect_job_from_redis(cache_key)
if isinstance(cached, dict):
self._set_local_active_detect_job_cache(
cached,
now_ts=time.time(),
fresh_ttl_seconds=local_cache_ttl_seconds,
stale_ttl_seconds=stale_cache_ttl_seconds,
)
return self._clone_active_detect_job_payload(cached)
refresh_lock_token = ""
except Exception as e:
logger.debug(f"获取活动任务缓存刷新锁失败: {e}")
refresh_lock_token = ""
tail_handoff_enabled = str(
os.getenv("DOMAINCHECK_TAIL_HANDOFF_ENABLED", "1") or "1"
).strip().lower() not in {"0", "false", "off", "no"}
try:
tail_handoff_active_max = max(
0,
int(
os.getenv(
"DOMAINCHECK_TAIL_HANDOFF_MAX_ACTIVE_ITEMS",
"128",
)
or 128
),
)
except (TypeError, ValueError):
tail_handoff_active_max = 128
try:
tail_handoff_pending_max = max(
0,
int(
os.getenv(
"DOMAINCHECK_TAIL_HANDOFF_MAX_PENDING_ITEMS",
"256",
)
or 256
),
)
except (TypeError, ValueError):
tail_handoff_pending_max = 256
try:
tail_handoff_min_pending = max(
1,
int(
os.getenv(
"DOMAINCHECK_TAIL_HANDOFF_MIN_PENDING_ITEMS",
"1",
)
or 1
),
)
except (TypeError, ValueError):
tail_handoff_min_pending = 1
try:
running_stall_seconds = max(
60,
int(
os.getenv(
"DOMAINCHECK_RUNNING_JOB_STALL_SECONDS",
"900",
)
or 900
),
)
except (TypeError, ValueError):
running_stall_seconds = 900
row = self.fetch_one(
"""
WITH tail_config AS (
SELECT
%s::boolean AS handoff_enabled,
%s::integer AS handoff_active_max,
%s::integer AS handoff_pending_max,
%s::integer AS handoff_min_pending,
%s::integer AS running_stall_seconds
),
candidate_jobs AS (
SELECT
job.id,
job.job_code,
job.task_mode,
job.status,
job.started_at,
job.created_at
FROM detect_jobs AS job
WHERE job.status IN ('pending', 'running')
ORDER BY
CASE WHEN job.status = 'running' THEN 0 ELSE 1 END ASC,
COALESCE(job.started_at, job.created_at) DESC,
job.id DESC
LIMIT 20
),
job_counts AS (
SELECT
job.id,
job.job_code,
job.task_mode,
job.status,
job.started_at,
job.created_at,
COALESCE(item_counts.items_pending, 0) AS items_pending,
COALESCE(item_counts.items_claimed, 0) AS items_claimed,
COALESCE(item_counts.items_running, 0) AS items_running,
COALESCE(item_counts.items_claimed, 0) + COALESCE(item_counts.items_running, 0) AS items_active,
item_counts.latest_unfinished_activity_at,
item_counts.latest_active_activity_at
FROM candidate_jobs AS job
LEFT JOIN LATERAL (
SELECT
count(*) FILTER (WHERE item.status = 'pending') AS items_pending,
count(*) FILTER (WHERE item.status = 'claimed') AS items_claimed,
count(*) FILTER (WHERE item.status = 'running') AS items_running,
max(COALESCE(item.updated_at, item.started_at, item.create_time)) AS latest_unfinished_activity_at,
max(COALESCE(item.updated_at, item.started_at, item.create_time))
FILTER (WHERE item.status IN ('claimed', 'running')) AS latest_active_activity_at
FROM detect_job_items AS item
WHERE item.job_id = job.id
AND item.status IN ('pending', 'claimed', 'running')
) AS item_counts ON TRUE
WHERE
COALESCE(item_counts.items_pending, 0) > 0
OR COALESCE(item_counts.items_claimed, 0) > 0
OR COALESCE(item_counts.items_running, 0) > 0
),
job_health AS (
SELECT
job.*,
CASE
WHEN job.status = 'running'
AND job.items_active > 0
AND cfg.running_stall_seconds > 0
AND COALESCE(
job.latest_active_activity_at,
job.latest_unfinished_activity_at,
job.started_at,
job.created_at
) <= (NOW() - (cfg.running_stall_seconds * INTERVAL '1 second'))
THEN TRUE
ELSE FALSE
END AS running_job_stalled
FROM job_counts AS job
CROSS JOIN tail_config AS cfg
),
tail_gate AS (
SELECT
COALESCE(
bool_or(
job.status = 'pending'
AND job.items_pending >= cfg.handoff_min_pending
),
FALSE
) AS has_handoff_pending,
COALESCE(
bool_or(
job.status = 'running'
AND NOT job.running_job_stalled
AND (
job.items_active > cfg.handoff_active_max
OR job.items_pending > cfg.handoff_pending_max
)
),
FALSE
) AS has_non_tail_running
FROM job_health AS job
CROSS JOIN tail_config AS cfg
)
SELECT
job.id,
job.job_code,
job.task_mode,
job.status,
job.items_pending,
job.items_claimed,
job.items_running,
0 AS items_completed,
0 AS items_failed,
job.latest_unfinished_activity_at AS latest_item_updated_at,
NULL::timestamp AS latest_item_created_at,
job.running_job_stalled,
CASE
WHEN cfg.handoff_enabled
AND gate.has_handoff_pending
AND NOT gate.has_non_tail_running
AND job.status = 'pending'
AND job.items_pending >= cfg.handoff_min_pending
THEN TRUE
ELSE FALSE
END AS tail_handoff_candidate,
CASE
WHEN cfg.handoff_enabled
AND gate.has_handoff_pending
AND NOT gate.has_non_tail_running
AND job.status = 'pending'
AND job.items_pending >= cfg.handoff_min_pending
THEN 'tail_handoff_pending'
WHEN job.running_job_stalled
THEN 'running_job_stalled'
WHEN job.items_claimed > 0 OR job.items_running > 0
THEN 'running_job_active'
WHEN job.items_pending > 0
THEN 'pending_job'
ELSE 'inactive'
END AS selection_reason
FROM job_health AS job
CROSS JOIN tail_gate AS gate
CROSS JOIN tail_config AS cfg
ORDER BY
CASE
WHEN cfg.handoff_enabled
AND gate.has_handoff_pending
AND NOT gate.has_non_tail_running
AND job.status = 'pending'
AND job.items_pending >= cfg.handoff_min_pending
THEN 0
WHEN (job.items_claimed > 0 OR job.items_running > 0)
AND NOT job.running_job_stalled
THEN 1
WHEN job.items_pending > 0 THEN 2
WHEN job.running_job_stalled THEN 3
ELSE 4
END ASC,
CASE
WHEN cfg.handoff_enabled
AND gate.has_handoff_pending
AND NOT gate.has_non_tail_running
AND job.status = 'pending'
AND job.items_pending >= cfg.handoff_min_pending
THEN job.items_pending
ELSE NULL
END DESC NULLS LAST,
CASE WHEN job.status = 'running' THEN 0 ELSE 1 END ASC,
COALESCE(job.started_at, job.created_at) DESC,
job.id DESC
LIMIT 1
""",
(
tail_handoff_enabled,
tail_handoff_active_max,
tail_handoff_pending_max,
tail_handoff_min_pending,
running_stall_seconds,
),
)
try:
if isinstance(row, dict):
cache_row = {}
for key, value in row.items():
if isinstance(value, datetime):
cache_row[key] = value.isoformat()
else:
cache_row[key] = value
self._set_local_active_detect_job_cache(
cache_row,
now_ts=time.time(),
fresh_ttl_seconds=local_cache_ttl_seconds,
stale_ttl_seconds=stale_cache_ttl_seconds,
)
if self.redis_client:
try:
self.redis_client.setex(
cache_key,
cache_ttl_seconds,
json.dumps(cache_row, ensure_ascii=True, default=str),
)
except Exception as e:
logger.debug(f"写入活动任务缓存失败: {e}")
return row
self._set_local_active_detect_job_cache(
row,
now_ts=time.time(),
fresh_ttl_seconds=local_cache_ttl_seconds,
stale_ttl_seconds=stale_cache_ttl_seconds,
)
return row
finally:
if self.redis_client and refresh_lock_token:
try:
current_owner = self.redis_client.get(cache_lock_key)
if current_owner and str(current_owner) == refresh_lock_token:
self.redis_client.delete(cache_lock_key)
except Exception:
pass
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'
RETURNING item.job_id
"""
)
updated_rows = cur.fetchall() or []
updated_count = len(updated_rows)
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, item.job_id
"""
)
updated_rows = cur.fetchall() or []
updated_item_ids = {int(row[0]) for row in updated_rows}
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 COALESCE(task_mode, '')
FROM detect_jobs
WHERE id = %s
""",
(job_id,),
)
task_mode_row = cur.fetchone()
task_mode = str((task_mode_row or [""])[0] or "").strip()
def _exists(query, *params):
cur.execute(query, params)
row = cur.fetchone()
return bool((row or [False])[0])
dispatch_active_exists = _exists(
"""
SELECT EXISTS (
SELECT 1
FROM detect_job_items AS item
WHERE item.job_id = %s
AND item.status IN ('claimed', 'running')
LIMIT 1
)
""",
job_id,
)
unprocessed_terminal_exists = False
if task_mode == 'domain_pipeline':
unprocessed_terminal_exists = _exists(
"""
SELECT EXISTS (
SELECT 1
FROM detect_job_items AS item
WHERE item.job_id = %s
AND item.status IN ('completed', 'blacklisted', 'failed')
AND COALESCE(item.step_code, '') <> ''
AND COALESCE(item.result_payload_json->>'controller_processed', 'false') <> 'true'
LIMIT 1
)
""",
job_id,
)
if dispatch_active_exists or unprocessed_terminal_exists:
cur.execute(
"""
UPDATE detect_jobs
SET status = 'running',
started_at = COALESCE(started_at, CURRENT_TIMESTAMP),
finished_at = NULL
WHERE id = %s
AND (
status <> 'running'
OR started_at IS NULL
OR finished_at IS NOT NULL
)
""",
(job_id,),
)
return
pending_exists = _exists(
"""
SELECT EXISTS (
SELECT 1
FROM detect_job_items AS item
WHERE item.job_id = %s
AND item.status = 'pending'
LIMIT 1
)
""",
job_id,
)
if pending_exists:
cur.execute(
"""
UPDATE detect_jobs
SET status = 'pending',
finished_at = NULL
WHERE id = %s
""",
(job_id,),
)
return
failed_exists = _exists(
"""
SELECT EXISTS (
SELECT 1
FROM detect_job_items AS item
WHERE item.job_id = %s
AND item.status = 'failed'
LIMIT 1
)
""",
job_id,
)
done_exists = _exists(
"""
SELECT EXISTS (
SELECT 1
FROM detect_job_items AS item
WHERE item.job_id = %s
AND item.status IN ('completed', 'blacklisted')
LIMIT 1
)
""",
job_id,
)
final_status = 'completed'
if failed_exists and done_exists:
final_status = 'partial_failed'
elif failed_exists:
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)
Reference in New Issue View Git Blame Copy Permalink