mdbq 3.10.7__py3-none-any.whl → 3.10.9__py3-none-any.whl

mdbq/mysql/s_query.py

@@ -1,14 +1,11 @@
 # -*- coding:utf-8 -*-
 import datetime
-import re
-import time
 import warnings
 import pymysql
-import numpy as np
 import pandas as pd
-import os
 from decimal import Decimal
 import logging
+from contextlib import closing
 
 warnings.filterwarnings('ignore')
 """
@@ -18,7 +15,20 @@ logger = logging.getLogger(__name__)
 
 
 class QueryDatas:
+    """
+    数据库查询工具类。
+    用于连接MySQL数据库，支持表结构检查、条件查询、数据导出为DataFrame、列名和类型获取等功能。
+    """
+
     def __init__(self, username: str, password: str, host: str, port: int, charset: str = 'utf8mb4'):
+        """
+        初始化数据库连接配置。
+        :param username: 数据库用户名
+        :param password: 数据库密码
+        :param host: 数据库主机
+        :param port: 数据库端口
+        :param charset: 字符集，默认utf8mb4
+        """
         self.username = username
         self.password = password
         self.host = host
@@ -32,175 +42,181 @@ class QueryDatas:
             'cursorclass': pymysql.cursors.DictCursor,
         }
 
-    def check_condition(self, db_name, table_name, condition):
-        """ 按指定条件查询数据库，并返回 """
-        if self.check_infos(db_name, table_name) == False:
-            return
-
+    def check_condition(self, db_name, table_name, condition, columns='更新时间'):
+        """
+        按指定条件查询数据库表，返回满足条件的指定字段数据。
+        :param db_name: 数据库名
+        :param table_name: 表名
+        :param condition: SQL条件字符串（不含WHERE）
+        :param columns: 查询字段字符串或以逗号分隔的字段名，默认'更新时间'
+        :return: 查询结果列表或None
+        """
+        if not self.check_infos(db_name, table_name):
+            return None
         self.config.update({'database': db_name})
-        connection = pymysql.connect(**self.config)  # 重新连接数据库
-        with connection.cursor() as cursor:
-            sql = f"SELECT 更新时间 FROM {table_name} WHERE {condition}"
-            # logger.info(sql)
-            cursor.execute(sql)
-            columns = cursor.fetchall()
-            return columns
+        try:
+            with closing(pymysql.connect(**self.config)) as connection:
+                with closing(connection.cursor()) as cursor:
+                    sql = f"SELECT {columns} FROM `{table_name}` WHERE {condition}"
+                    logger.debug(f"check_condition SQL: {sql}")
+                    cursor.execute(sql)
+                    result = cursor.fetchall()
+                    return result
+        except Exception as e:
+            logger.error(f"check_condition error: {e}")
+            return None
 
-    def data_to_df(self, db_name, table_name, start_date, end_date, projection: dict = None):
+    def data_to_df(self, db_name, table_name, start_date, end_date, projection: dict = None, limit: int = None):
         """
-        从数据库表获取数据到DataFrame，支持列筛选和日期范围过滤
-        Args:
-            db_name: 数据库名
-            table_name: 表名
-            start_date: 起始日期（包含）
-            end_date: 结束日期（包含）
-            projection: 列筛选字典，e.g. {'日期': 1, '场景名字': 1}
+        从数据库表获取数据到DataFrame，支持列筛选、日期范围过滤和行数限制。
+        :param db_name: 数据库名
+        :param table_name: 表名
+        :param start_date: 起始日期（包含）
+        :param end_date: 结束日期（包含）
+        :param projection: 列筛选字典，e.g. {'日期': 1, '场景名字': 1}
+        :param limit: 限制返回的最大行数
+        :return: 查询结果的DataFrame
         """
-        # 初始化默认参数
         projection = projection or {}
         df = pd.DataFrame()
-        # 日期处理
-        start_date = pd.to_datetime(start_date or '1970-01-01').strftime('%Y-%m-%d')
-        end_date = pd.to_datetime(end_date or datetime.datetime.today()).strftime('%Y-%m-%d')
-
-        # 前置检查
+        try:
+            start_date = pd.to_datetime(start_date or '1970-01-01').strftime('%Y-%m-%d')
+            end_date = pd.to_datetime(end_date or datetime.datetime.today()).strftime('%Y-%m-%d')
+        except Exception as e:
+            logger.error(f"日期格式错误: {e}")
+            return df
         if not self.check_infos(db_name, table_name):
             return df
-
-        # 配置数据库连接
         self.config['database'] = db_name
-        connection = None
-
         try:
-            connection = pymysql.connect(**self.config)
-            with connection.cursor() as cursor:
-                # 获取表结构（排除id列）
-                cursor.execute(
-                    """SELECT COLUMN_NAME 
-                    FROM information_schema.columns 
-                    WHERE table_schema = %s AND table_name = %s""",
-                    (db_name, table_name)
-                )
-                cols_exist = {col['COLUMN_NAME'] for col in cursor.fetchall()} - {'id'}
-
-                # 处理列选择
-                selected_columns = []
-                if projection:
-                    selected_columns = [k for k, v in projection.items() if v and k in cols_exist]
+            with closing(pymysql.connect(**self.config)) as connection:
+                with closing(connection.cursor()) as cursor:
+                    cursor.execute(
+                        """SELECT COLUMN_NAME FROM information_schema.columns WHERE table_schema = %s AND table_name = %s""",
+                        (db_name, table_name)
+                    )
+                    cols_exist = {col['COLUMN_NAME'] for col in cursor.fetchall()} - {'id'}
+                    if projection:
+                        selected_columns = [k for k, v in projection.items() if v and k in cols_exist]
+                        if not selected_columns:
+                            logger.info("Warning: Projection 参数不匹配任何数据库字段")
+                            return df
+                    else:
+                        selected_columns = list(cols_exist)
                     if not selected_columns:
-                        logger.info("Warning: Projection 参数不匹配任何数据库字段")
+                        logger.info("未找到可用字段")
                         return df
-                else:
-                    selected_columns = list(cols_exist)
-                # 构建基础SQL
-                quoted_columns = [f'`{col}`' for col in selected_columns]
-                base_sql = f"SELECT {', '.join(quoted_columns)} FROM `{db_name}`.`{table_name}`"
-
-                # 添加日期条件
-                if '日期' in cols_exist:
-                    base_sql += f" WHERE 日期 BETWEEN '{start_date}' AND '{end_date}'"
-
-                # 执行查询
-                cursor.execute(base_sql)
-                result = cursor.fetchall()
-
-                # 处理结果集
-                if result:
-                    df = pd.DataFrame(result, columns=[desc[0] for desc in cursor.description])
-                    # 类型转换优化
-                    decimal_cols = [col for col in df.columns if df[col].apply(lambda x: isinstance(x, Decimal)).any()]
-                    df[decimal_cols] = df[decimal_cols].astype(float)
-
+                    quoted_columns = [f'`{col}`' for col in selected_columns]
+                    base_sql = f"SELECT {', '.join(quoted_columns)} FROM `{db_name}`.`{table_name}`"
+                    params = []
+                    if '日期' in cols_exist:
+                        base_sql += f" WHERE 日期 BETWEEN %s AND %s"
+                        params.extend([start_date, end_date])
+                    if limit is not None and isinstance(limit, int) and limit > 0:
+                        base_sql += f" LIMIT %s"
+                        params.append(limit)
+                    logger.debug(f"data_to_df SQL: {base_sql}, params: {params}")
+                    cursor.execute(base_sql, tuple(params))
+                    result = cursor.fetchall()
+                    if result:
+                        df = pd.DataFrame(result)
+                        for col in df.columns:
+                            if df[col].apply(lambda x: isinstance(x, Decimal)).any():
+                                df[col] = df[col].astype(float)
         except Exception as e:
-            logger.error(f"Database operation failed: {str(e)}")
-        finally:
-            if connection:
-                connection.close()
-
+            logger.error(f"data_to_df error: {e}")
         return df
 
-    def columns_to_list(self, db_name, table_name,  columns_name) -> list:
+    def columns_to_list(self, db_name, table_name, columns_name, where: str = None) -> list:
         """
-        获取数据表的指定列, 返回列表
-        [{'视频bv号': 'BV1Dm4y1S7BU', '下载进度': 1}, {'视频bv号': 'BV1ov411c7US', '下载进度': 1}]
+        获取数据表的指定列, 支持where条件筛选, 返回列表字典。
+        :param db_name: 数据库名
+        :param table_name: 表名
+        :param columns_name: 需要获取的列名列表
+        :param where: 可选，SQL条件字符串（不含WHERE）
+        :return: [{列1:值, 列2:值, ...}, ...]
         """
-        if self.check_infos(db_name, table_name) == False:  # 检查传入的数据库和数据表是否存在
+        if not self.check_infos(db_name, table_name):
             return []
-
         self.config.update({'database': db_name})
-        connection = pymysql.connect(**self.config)  # 重新连接数据库
-        with connection.cursor() as cursor:
-            # 3. 获取数据表的所有列信息
-            sql = 'SELECT COLUMN_NAME FROM information_schema.columns WHERE table_schema = %s AND table_name = %s'
-            cursor.execute(sql, (db_name, {table_name}))
-            columns = cursor.fetchall()
-            cols_exist = [col['COLUMN_NAME'] for col in columns]  # 数据表的所有列, 返回 list
-            columns_name = [item for item in columns_name if item in cols_exist]
-            if len(columns_name) == 0:
-                return []
-            columns_in = ', '.join(columns_name)
-            sql = (f"SELECT {columns_in} FROM {db_name}.{table_name} ")
-            cursor.execute(sql)
-            column_values = cursor.fetchall()  # 返回指定列，结果是[dict, dict, dict, ...]
-            # column_values = [item[column_name] for item in column_values]  # 提取字典的值, 组成列表
-        connection.close()
-        return column_values
+        try:
+            with closing(pymysql.connect(**self.config)) as connection:
+                with closing(connection.cursor()) as cursor:
+                    sql = 'SELECT COLUMN_NAME FROM information_schema.columns WHERE table_schema = %s AND table_name = %s'
+                    cursor.execute(sql, (db_name, table_name))
+                    cols_exist = [col['COLUMN_NAME'] for col in cursor.fetchall()]
+                    columns_name = [item for item in columns_name if item in cols_exist]
+                    if not columns_name:
+                        logger.info("columns_to_list: 未找到匹配的列名")
+                        return []
+                    columns_in = ', '.join([f'`{col}`' for col in columns_name])
+                    sql = f"SELECT {columns_in} FROM `{db_name}`.`{table_name}`"
+                    if where:
+                        sql += f" WHERE {where}"
+                    logger.debug(f"columns_to_list SQL: {sql}")
+                    cursor.execute(sql)
+                    column_values = cursor.fetchall()
+            return column_values
+        except Exception as e:
+            logger.error(f"columns_to_list error: {e}")
+            return []
 
-    def dtypes_to_list(self, db_name, table_name) -> list:
+    def dtypes_to_list(self, db_name, table_name, columns_name=None) -> list:
         """
-        获取数据表的指定列, 返回列表
-        [{'视频bv号': 'BV1Dm4y1S7BU', '下载进度': 1}, {'视频bv号': 'BV1ov411c7US', '下载进度': 1}]
+        获取数据表的列名和类型, 支持只返回部分字段类型。
+        :param db_name: 数据库名
+        :param table_name: 表名
+        :param columns_name: 可选，字段名列表，仅返回这些字段的类型
+        :return: [{'COLUMN_NAME': ..., 'COLUMN_TYPE': ...}, ...]
         """
-        if self.check_infos(db_name, table_name) == False:  # 检查传入的数据库和数据表是否存在
+        if not self.check_infos(db_name, table_name):
             return []
-
         self.config.update({'database': db_name})
-        connection = pymysql.connect(**self.config)  # 重新连接数据库
-        with connection.cursor() as cursor:
-            # 3. 获取数据表的所有列信息
-            sql = 'SELECT COLUMN_NAME, COLUMN_TYPE FROM information_schema.columns WHERE table_schema = %s AND table_name = %s'
-            cursor.execute(sql, (db_name, {table_name}))
-            column_name_and_type = cursor.fetchall()
-        connection.close()
-        return column_name_and_type
+        try:
+            with closing(pymysql.connect(**self.config)) as connection:
+                with closing(connection.cursor()) as cursor:
+                    sql = 'SELECT COLUMN_NAME, COLUMN_TYPE FROM information_schema.columns WHERE table_schema = %s AND table_name = %s'
+                    cursor.execute(sql, (db_name, table_name))
+                    column_name_and_type = cursor.fetchall()
+                    if columns_name:
+                        columns_name = set(columns_name)
+                        column_name_and_type = [row for row in column_name_and_type if row['COLUMN_NAME'] in columns_name]
+            return column_name_and_type
+        except Exception as e:
+            logger.error(f"dtypes_to_list error: {e}")
+            return []
 
     def check_infos(self, db_name, table_name) -> bool:
-        """ 检查数据库、数据表是否存在 """
-        connection = pymysql.connect(**self.config)  # 连接数据库
+        """
+        检查数据库和数据表是否存在。
+        :param db_name: 数据库名
+        :param table_name: 表名
+        :return: 存在返回True，否则False
+        """
         try:
-            with connection.cursor() as cursor:
-                # 1. 检查数据库是否存在
-                cursor.execute(f"SHOW DATABASES LIKE '{db_name}'")  # 检查数据库是否存在
-                database_exists = cursor.fetchone()
-                if not database_exists:
-                    logger.info(f"Database <{db_name}>: 数据库不存在")
-                    return False
-        finally:
-            connection.close()  # 这里要断开连接
-
-        self.config.update({'database': db_name})  # 添加更新 config 字段
-        connection = pymysql.connect(**self.config)  # 重新连接数据库
+            with closing(pymysql.connect(**self.config)) as connection:
+                with closing(connection.cursor()) as cursor:
+                    cursor.execute(f"SHOW DATABASES LIKE %s", (db_name,))
+                    database_exists = cursor.fetchone()
+                    if not database_exists:
+                        logger.info(f"Database <{db_name}>: 数据库不存在")
+                        return False
+        except Exception as e:
+            logger.error(f"check_infos-db error: {e}")
+            return False
+        self.config.update({'database': db_name})
         try:
-            with connection.cursor() as cursor:
-                # 2. 查询表是否存在
-                sql = f"SHOW TABLES LIKE '{table_name}'"
-                cursor.execute(sql)
-                if not cursor.fetchone():
-                    logger.info(f'{db_name} -> <{table_name}>: 表不存在')
-                    return False
-                return True
+            with closing(pymysql.connect(**self.config)) as connection:
+                with closing(connection.cursor()) as cursor:
+                    cursor.execute(f"SHOW TABLES LIKE %s", (table_name,))
+                    if not cursor.fetchone():
+                        logger.info(f'{db_name} -> <{table_name}>: 表不存在')
+                        return False
+                    return True
         except Exception as e:
-            logger.error(e)
+            logger.error(f"check_infos-table error: {e}")
             return False
-        finally:
-            connection.close()  # 断开连接
 
 
 if __name__ == '__main__':
-    conf = ConfigTxt()
-    data = conf.config_datas['Windows']['xigua_lx']['mysql']['remoto']
-    username, password, host, port = data['username'], data['password'], data['host'], data['port']
-
-    q = QueryDatas(username, password, host, port)
-    res = q.columns_to_list(db_name='视频数据', table_name='bilibili视频', columns_name=['视频bv号', '下载进度'])
-    logger.info(res)
+    pass