mdbq 3.10.10__py3-none-any.whl → 3.10.12__py3-none-any.whl

mdbq/__version__.py

@@ -1 +1 @@
-VERSION = '3.10.10'
+VERSION = '3.10.12'

mdbq/mysql/deduplicator.py

@@ -12,6 +12,7 @@ import threading
 import concurrent.futures
 from collections import defaultdict
 import sys
+from datetime import datetime
 
 
 warnings.filterwarnings('ignore')
@@ -77,15 +78,18 @@ class MySQLDeduplicator:
             date_range: Optional[List[str]] = None,
             recent_month: Optional[int] = None,
             date_column: str = '日期',
-            exclude_columns: Optional[List[str]] = None
+            exclude_columns: Optional[List[str]] = None,
+            exclude_databases: Optional[List[str]] = None,
+            exclude_tables: Optional[Dict[str, List[str]]] = None
     ) -> None:
         """
         初始化去重处理器
-        新增参数:
         :param date_range: 指定去重的日期区间 [start_date, end_date]，格式'YYYY-MM-DD'
         :param recent_month: 最近N个月的数据去重（与date_range互斥，优先生效）
         :param date_column: 时间列名，默认为'日期'
         :param exclude_columns: 去重时排除的列名列表，默认为['id', '更新时间']
+        :param exclude_databases: 排除的数据库名列表
+        :param exclude_tables: 排除的表名字典 {数据库名: [表名, ...]}
         """
         # 连接池状态标志
         self._closed = False
@@ -127,11 +131,28 @@ class MySQLDeduplicator:
             self.exclude_columns = list(default_exclude | {'更新时间'})
         else:
             self.exclude_columns = list(set(exclude_columns) | default_exclude)
-        # 解析时间范围
+        # 解析时间范围并智能校正date_range
         if self.date_range and len(self.date_range) == 2:
-            self._dedup_start_date, self._dedup_end_date = self.date_range
+            try:
+                start, end = self.date_range
+                start_dt = datetime.strptime(start, "%Y-%m-%d")
+                end_dt = datetime.strptime(end, "%Y-%m-%d")
+                if start_dt > end_dt:
+                    logger.warning(
+                        "date_range顺序不正确，自动交换开始和结束日期。",
+                        {"start": start, "end": end}
+                    )
+                    start_dt, end_dt = end_dt, start_dt
+                self._dedup_start_date = start_dt.strftime("%Y-%m-%d")
+                self._dedup_end_date = end_dt.strftime("%Y-%m-%d")
+            except Exception as e:
+                logger.error(
+                    "date_range参数格式错误，应为['YYYY-MM-DD', 'YYYY-MM-DD']，已忽略时间范围。",
+                    {"date_range": self.date_range, "error": str(e)}
+                )
+                self._dedup_start_date = None
+                self._dedup_end_date = None
         elif self.recent_month:
-            from datetime import datetime, timedelta
             today = datetime.today()
             month = today.month - self.recent_month
             year = today.year
@@ -148,8 +169,19 @@ class MySQLDeduplicator:
         # 系统数据库列表
         self.SYSTEM_DATABASES = {'information_schema', 'mysql', 'performance_schema', 'sys'}
 
+        # 排除数据库和表的逻辑
+        self.exclude_databases = set([db.lower() for db in exclude_databases]) if exclude_databases else set()
+        self.exclude_tables = {k.lower(): set([t.lower() for t in v]) for k, v in (exclude_tables or {}).items()}
+
     def _get_connection(self) -> pymysql.connections.Connection:
-        """从连接池获取连接"""
+        """
+        从连接池获取一个数据库连接。
+        
+        Returns:
+            pymysql.connections.Connection: 数据库连接对象。
+        Raises:
+            ConnectionError: 如果连接池已关闭或获取连接失败。
+        """
         if self._closed:
             logger.error('尝试获取连接但连接池已关闭')
             raise ConnectionError("连接池已关闭")
@@ -163,7 +195,17 @@ class MySQLDeduplicator:
 
     @staticmethod
     def _retry_on_failure(func: Any) -> Any:
-        """重试装饰器"""
+        """
+        装饰器：为数据库操作方法提供自动重试机制。
+        仅捕获pymysql的连接相关异常，重试指定次数后抛出最后一次异常。
+        
+        Args:
+            func (Any): 被装饰的函数。
+        Returns:
+            Any: 被装饰函数的返回值。
+        Raises:
+            Exception: 多次重试后依然失败时抛出。
+        """
         @wraps(func)
         def wrapper(self, *args, **kwargs):
             last_exception = None
@@ -192,21 +234,31 @@ class MySQLDeduplicator:
 
     @_retry_on_failure
     def _get_databases(self) -> List[str]:
-        """获取所有非系统数据库列表"""
+        """
+        获取所有非系统数据库列表，排除exclude_databases。
+        
+        Returns:
+            List[str]: 数据库名列表。
+        """
         sql = "SHOW DATABASES"
-
         with self._get_connection() as conn:
             with conn.cursor() as cursor:
                 cursor.execute(sql)
                 all_dbs = [row['Database'] for row in cursor.fetchall()]
-
-                if self.skip_system_dbs:
-                    return [db for db in all_dbs if db.lower() not in self.SYSTEM_DATABASES]
-                return all_dbs
+                # 排除系统库和用户指定的库
+                filtered = [db for db in all_dbs if db.lower() not in self.SYSTEM_DATABASES and db.lower() not in self.exclude_databases] if self.skip_system_dbs else [db for db in all_dbs if db.lower() not in self.exclude_databases]
+                return filtered
 
     @_retry_on_failure
     def _get_tables(self, database: str) -> List[str]:
-        """获取指定数据库的所有表"""
+        """
+        获取指定数据库的所有表名。
+        
+        Args:
+            database (str): 数据库名。
+        Returns:
+            List[str]: 表名列表。
+        """
         sql = "SHOW TABLES"
 
         with self._get_connection() as conn:
@@ -217,7 +269,15 @@ class MySQLDeduplicator:
 
     @_retry_on_failure
     def _get_table_columns(self, database: str, table: str) -> List[str]:
-        """获取表的列名(排除主键列)"""
+        """
+        获取指定表的所有列名（排除主键列）。
+        
+        Args:
+            database (str): 数据库名。
+            table (str): 表名。
+        Returns:
+            List[str]: 列名列表。
+        """
         sql = """
         SELECT COLUMN_NAME 
         FROM INFORMATION_SCHEMA.COLUMNS 
@@ -232,7 +292,15 @@ class MySQLDeduplicator:
                         if row['COLUMN_NAME'].lower() != self.primary_key.lower()]
 
     def _acquire_table_lock(self, database: str, table: str) -> bool:
-        """获取表处理锁，防止并发处理同一张表"""
+        """
+        获取表处理锁，防止并发处理同一张表。
+        
+        Args:
+            database (str): 数据库名。
+            table (str): 表名。
+        Returns:
+            bool: 是否成功获取锁。
+        """
         key = f"{database}.{table}"
 
         with self._lock:
@@ -243,7 +311,13 @@ class MySQLDeduplicator:
             return True
 
     def _release_table_lock(self, database: str, table: str) -> None:
-        """释放表处理锁"""
+        """
+        释放表处理锁。
+        
+        Args:
+            database (str): 数据库名。
+            table (str): 表名。
+        """
         key = f"{database}.{table}"
 
         with self._lock:
@@ -258,13 +332,15 @@ class MySQLDeduplicator:
             dry_run: bool = False
     ) -> Tuple[int, int]:
         """
-        执行单表去重
-
-        :param database: 数据库名
-        :param table: 表名
-        :param columns: 用于去重的列(为None时使用所有列)
-        :param dry_run: 是否模拟运行(只统计不实际删除)
-        :return: (重复行数, 删除行数)
+        执行单表去重。
+        
+        Args:
+            database (str): 数据库名。
+            table (str): 表名。
+            columns (Optional[List[str]]): 用于去重的列名列表（为None时使用所有列）。
+            dry_run (bool): 是否为模拟运行（只统计不实际删除）。
+        Returns:
+            Tuple[int, int]: (重复组数, 实际删除行数)。
         """
         if not self._acquire_table_lock(database, table):
             return (0, 0)
@@ -385,14 +461,19 @@ class MySQLDeduplicator:
             dry_run: bool = False
     ) -> Tuple[int, int]:
         """
-        对指定表进行去重
-
-        :param database: 数据库名
-        :param table: 表名
-        :param columns: 用于去重的列(为None时使用所有列)
-        :param dry_run: 是否模拟运行(只统计不实际删除)
-        :return: (重复行数, 删除行数)
+        对指定表进行去重。
+        
+        Args:
+            database (str): 数据库名。
+            table (str): 表名。
+            columns (Optional[List[str]]): 用于去重的列名列表（为None时使用所有列）。
+            dry_run (bool): 是否为模拟运行（只统计不实际删除）。
+        Returns:
+            Tuple[int, int]: (重复组数, 实际删除行数)。
         """
+        if database.lower() in self.exclude_tables and table.lower() in self.exclude_tables[database.lower()]:
+            logger.info('表被排除', {"库": database, "表": table, "操作": "跳过"})
+            return (0, 0)
         try:
             if not self._check_table_exists(database, table):
                 logger.warning('表不存在', {"库": database, "表": table, "warning": "跳过"})
@@ -414,14 +495,16 @@ class MySQLDeduplicator:
             parallel: bool = False
     ) -> Dict[str, Tuple[int, int]]:
         """
-        对指定数据库的所有表进行去重
-
-        :param database: 数据库名
-        :param tables: 要处理的表列表(为None时处理所有表)
-        :param columns_map: 各表使用的去重列 {表名: [列名]}
-        :param dry_run: 是否模拟运行
-        :param parallel: 是否并行处理
-        :return: 字典 {表名: (重复行数, 删除行数)}
+        对指定数据库的所有表进行去重。
+        
+        Args:
+            database (str): 数据库名。
+            tables (Optional[List[str]]): 要处理的表列表（为None时处理所有表）。
+            columns_map (Optional[Dict[str, List[str]]]): 各表使用的去重列 {表名: [列名]}。
+            dry_run (bool): 是否为模拟运行。
+            parallel (bool): 是否并行处理。
+        Returns:
+            Dict[str, Tuple[int, int]]: {表名: (重复组数, 实际删除行数)}。
         """
         results = {}
         try:
@@ -429,6 +512,8 @@ class MySQLDeduplicator:
                 logger.warning('数据库不存在', {"库": database})
                 return results
             target_tables = tables or self._get_tables(database)
+            exclude_tbls = self.exclude_tables.get(database.lower(), set())
+            target_tables = [t for t in target_tables if t.lower() not in exclude_tbls]
             logger.debug('获取目标表', {'库': database, 'tables': target_tables})
             if not target_tables:
                 logger.info('数据库中没有表', {"库": database, "操作": "跳过"})
@@ -482,18 +567,21 @@ class MySQLDeduplicator:
             parallel: bool = False
     ) -> Dict[str, Dict[str, Tuple[int, int]]]:
         """
-        对所有数据库进行去重
-
-        :param databases: 要处理的数据库列表(为None时处理所有非系统数据库)
-        :param tables_map: 各数据库要处理的表 {数据库名: [表名]}
-        :param columns_map: 各表使用的去重列 {数据库名: {表名: [列名]}}
-        :param dry_run: 是否模拟运行
-        :param parallel: 是否并行处理
-        :return: 嵌套字典 {数据库名: {表名: (重复行数, 删除行数)}}
+        对所有数据库进行去重。
+        
+        Args:
+            databases (Optional[List[str]]): 要处理的数据库列表。如果为 None，则处理所有非系统数据库。
+            tables_map (Optional[Dict[str, List[str]]]): 指定每个数据库要处理的表，格式为 {数据库名: [表名, ...]}。如果为 None，则处理所有表。
+            columns_map (Optional[Dict[str, Dict[str, List[str]]]]): 指定每个表去重时使用的列，格式为 {数据库名: {表名: [列名, ...]}}。如果为 None，则使用所有列。
+            dry_run (bool): 是否为模拟运行模式。为 True 时只统计重复行数，不实际删除。
+            parallel (bool): 是否并行处理多个数据库。为 True 时使用线程池并发处理。
+        Returns:
+            Dict[str, Dict[str, Tuple[int, int]]]: 嵌套字典，格式为 {数据库名: {表名: (重复组数, 实际删除行数)}}。
         """
-        all_results = defaultdict(dict)
+        all_results: Dict[str, Dict[str, Tuple[int, int]]] = defaultdict(dict)
         try:
-            target_dbs = databases or self._get_databases()
+            target_dbs: List[str] = databases or self._get_databases()
+            target_dbs = [db for db in target_dbs if db.lower() not in self.exclude_databases]
             logger.debug('获取目标数据库', {'databases': target_dbs})
             if not target_dbs:
                 logger.warning('没有可处理的数据库')
@@ -504,7 +592,7 @@ class MySQLDeduplicator:
                 with concurrent.futures.ThreadPoolExecutor(
                         max_workers=self.max_workers
                 ) as executor:
-                    futures = {}
+                    futures: Dict[concurrent.futures.Future, str] = {}
                     for db in target_dbs:
                         tables = tables_map.get(db) if tables_map else None
                         db_columns_map = columns_map.get(db) if columns_map else None
@@ -545,7 +633,14 @@ class MySQLDeduplicator:
 
     @_retry_on_failure
     def _check_database_exists(self, database: str) -> bool:
-        """检查数据库是否存在"""
+        """
+        检查数据库是否存在。
+        
+        Args:
+            database (str): 数据库名。
+        Returns:
+            bool: 数据库是否存在。
+        """
         sql = "SELECT SCHEMA_NAME FROM INFORMATION_SCHEMA.SCHEMATA WHERE SCHEMA_NAME = %s"
 
         with self._get_connection() as conn:
@@ -555,7 +650,15 @@ class MySQLDeduplicator:
 
     @_retry_on_failure
     def _check_table_exists(self, database: str, table: str) -> bool:
-        """检查表是否存在"""
+        """
+        检查表是否存在。
+        
+        Args:
+            database (str): 数据库名。
+            table (str): 表名。
+        Returns:
+            bool: 表是否存在。
+        """
         sql = """
         SELECT TABLE_NAME 
         FROM INFORMATION_SCHEMA.TABLES 
@@ -568,7 +671,12 @@ class MySQLDeduplicator:
                 return bool(cursor.fetchone())
 
     def close(self) -> None:
-        """关闭连接池"""
+        """
+        关闭连接池。
+        
+        Returns:
+            None
+        """
         try:
             if hasattr(self, 'pool') and self.pool and not self._closed:
                 self.pool.close()
@@ -580,9 +688,25 @@ class MySQLDeduplicator:
             logger.error(f"关闭连接池时出错", {'error_type': type(e).__name__, 'error': str(e)})
 
     def __enter__(self) -> 'MySQLDeduplicator':
+        """
+        上下文管理器进入方法。
+        
+        Returns:
+            MySQLDeduplicator: 实例自身。
+        """
         return self
 
     def __exit__(self, exc_type: Optional[type], exc_val: Optional[BaseException], exc_tb: Optional[Any]) -> None:
+        """
+        上下文管理器退出方法，自动关闭连接池。
+        
+        Args:
+            exc_type (Optional[type]): 异常类型。
+            exc_val (Optional[BaseException]): 异常值。
+            exc_tb (Optional[Any]): 异常追踪。
+        Returns:
+            None
+        """
         self.close()