getDomain/domainCheck/app/utils/database.py

# -*- coding: UTF-8 -*-
'''
@Project :domainScanDemo
@File    :database.py
@IDE     :PyCharm
@Author  :梦伴
@Date    :2026/4/9 0:07
@explain : 数据库操作类
'''

import psycopg2
import json
import redis
import threading
from loguru import logger
from app.config import config
from app.utils.status_codes import (
    DETECT_STATUS_BLACKLISTED,
    DETECT_STATUS_COMPLETED,
    DETECT_STATUS_FAILED,
    DETECT_STATUS_PENDING,
    REGISTER_STATUS_AVAILABLE,
    REGISTER_STATUS_REGISTERED,
    REVIEW_STATUS_PENDING,
    THIRD_PARTY_STATUS_DONE,
)


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 = config.DB_POOL_SIZE  # 连接池大小
        self.pool_lock = threading.Lock()
        
        # 初始化连接池
        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:
            for i in range(self.pool_size):
                conn = psycopg2.connect(
                    host=self.host,
                    port=self.port,
                    database=self.database,
                    user=self.user,
                    password=self.password
                )
                self.connection_pool.append(conn)
            logger.info(f"数据库连接池初始化成功，大小: {self.pool_size}")
        except Exception as e:
            logger.error(f"初始化数据库连接池失败: {e}")
    
    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.warning(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:
            with self.pool_lock:
                if not self.connection_pool:
                    # 连接池为空，尝试重新初始化
                    self._init_connection_pool()
                
                if self.connection_pool:
                    # 从连接池获取连接
                    conn = self.connection_pool.pop()
                    # 检查连接是否有效
                    if conn and not conn.closed:
                        try:
                            # 测试连接
                            cur = conn.cursor()
                            cur.execute("SELECT 1")
                            cur.fetchone()
                            cur.close()
                            logger.debug(f"线程 {thread_id} 从连接池获取连接成功")
                            return conn, conn.cursor()
                        except:
                            # 连接无效，关闭并重新获取
                            try:
                                conn.close()
                            except:
                                pass
                            if self.connection_pool:
                                conn = self.connection_pool.pop()
                                if conn and not conn.closed:
                                    return conn, conn.cursor()
                
                # 连接池为空或所有连接都无效，创建新连接
                logger.warning(f"连接池为空，线程 {thread_id} 创建新连接")
                conn = psycopg2.connect(
                    host=self.host,
                    port=self.port,
                    database=self.database,
                    user=self.user,
                    password=self.password
                )
                return conn, conn.cursor()
        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_lock:
                    if len(self.connection_pool) < self.pool_size:
                        self.connection_pool.append(conn)
                        logger.debug("连接已放回连接池")
                    else:
                        # 连接池已满，关闭连接
                        conn.close()
                        logger.debug("连接池已满，关闭连接")
        except Exception as e:
            logger.error(f"关闭数据库连接失败: {e}")
            try:
                if conn and not conn.closed:
                    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 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_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 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_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_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_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 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)