4310 lines
163 KiB
Python
4310 lines
163 KiB
Python
# -*- 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)
|