# -*- 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)