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