""" 数据查询引擎模块 ================ 本模块是BigDataTool的核心查询引擎，负责Cassandra数据库的高级查询功能。 核心功能： 1. 单表查询：标准的Cassandra CQL查询执行 2. 分表查询：基于TWCS策略的时间分表查询 3. 多主键查询：支持复合主键的复杂查询条件 4. 混合查询：生产环境分表+测试环境单表的组合查询 查询类型支持： - 单主键查询：WHERE key IN (val1, val2, val3) - 复合主键查询：WHERE (key1='val1' AND key2='val2') OR (key1='val3' AND key2='val4') - 分表查询：自动计算分表名称并并行查询多张表 - 字段过滤：支持指定查询字段和排除字段 分表查询特性： - 时间戳提取：从Key中智能提取时间戳信息 - 分表计算：基于TWCS策略计算目标分表 - 并行查询：同时查询多张分表以提高性能 - 错误容错：单个分表查询失败不影响整体结果 性能优化： - 查询时间监控：记录每个查询的执行时间 - 批量处理：支持大批量Key的高效查询 - 连接复用：优化数据库连接的使用 - 内存管理：大结果集的内存友好处理 作者：BigDataTool项目组 更新时间：2024年8月 """ import time import logging from .sharding import ShardingCalculator logger = logging.getLogger(__name__) def execute_query(session, table, keys, fields, values, exclude_fields=None): """ 执行Cassandra数据库查询，支持单主键和复合主键查询 本函数是查询引擎的核心，能够智能处理不同类型的主键查询： - 单主键：生成 WHERE key IN (val1, val2, val3) 查询 - 复合主键：生成 WHERE (key1='val1' AND key2='val2') OR ... 查询 Args: session: Cassandra数据库会话对象 table (str): 目标表名 keys (list): 主键字段名列表，如 ['id'] 或 ['docid', 'id'] fields (list): 要查询的字段列表，空列表表示查询所有字段 values (list): 查询值列表，复合主键值用逗号分隔 exclude_fields (list, optional): 要排除的字段列表 Returns: list: 查询结果列表，每个元素是一个Row对象 查询示例： # 单主键查询 execute_query(session, 'users', ['id'], ['name', 'email'], ['1', '2', '3']) # 生成SQL: SELECT name, email FROM users WHERE id IN ('1', '2', '3') # 复合主键查询 execute_query(session, 'orders', ['user_id', 'order_id'], ['*'], ['1,100', '2,200']) # 生成SQL: SELECT * FROM orders WHERE (user_id='1' AND order_id='100') OR (user_id='2' AND order_id='200') 错误处理： - 参数验证：检查keys和values是否为空 - SQL注入防护：对查询值进行适当转义 - 异常捕获：数据库错误时返回空列表 - 日志记录：记录查询SQL和执行统计 """ try: # 参数验证 if not keys or len(keys) == 0: logger.error("Keys参数为空，无法构建查询") return [] if not values or len(values) == 0: logger.error("Values参数为空，无法构建查询") return [] # 构建查询条件 if len(keys) == 1: # 单主键查询（保持原有逻辑） quoted_values = [f"'{value}'" for value in values] query_conditions = f"{keys[0]} IN ({', '.join(quoted_values)})" else: # 复合主键查询 conditions = [] for value in values: # 检查value是否包含复合主键分隔符 if isinstance(value, str) and ',' in value: # 解析复合主键值 key_values = [v.strip() for v in value.split(',')] if len(key_values) == len(keys): # 构建单个复合主键条件: (key1='val1' AND key2='val2') key_conditions = [] for i, (key, val) in enumerate(zip(keys, key_values)): key_conditions.append(f"{key} = '{val}'") conditions.append(f"({' AND '.join(key_conditions)})") else: logger.warning(f"复合主键值 '{value}' 的字段数量({len(key_values)})与主键字段数量({len(keys)})不匹配") # 将其作为第一个主键的值处理 conditions.append(f"{keys[0]} = '{value}'") else: # 单值，作为第一个主键的值处理 conditions.append(f"{keys[0]} = '{value}'") if conditions: query_conditions = ' OR '.join(conditions) else: logger.error("无法构建有效的查询条件") return [] # 确定要查询的字段 if fields: fields_str = ", ".join(fields) else: fields_str = "*" query_sql = f"SELECT {fields_str} FROM {table} WHERE {query_conditions};" # 记录查询SQL日志 logger.info(f"执行查询SQL: {query_sql}") if len(keys) > 1: logger.info(f"复合主键查询参数: 表={table}, 主键字段={keys}, 字段={fields_str}, Key数量={len(values)}") else: logger.info(f"单主键查询参数: 表={table}, 主键字段={keys[0]}, 字段={fields_str}, Key数量={len(values)}") # 执行查询 start_time = time.time() result = session.execute(query_sql) execution_time = time.time() - start_time result_list = list(result) if result else [] logger.info(f"查询完成: 执行时间={execution_time:.3f}秒, 返回记录数={len(result_list)}") return result_list except Exception as e: logger.error(f"查询执行失败: SQL={query_sql if 'query_sql' in locals() else 'N/A'}, 错误={str(e)}") return [] def execute_sharding_query(session, shard_mapping, keys, fields, exclude_fields=None): """ 执行分表查询 :param session: Cassandra会话 :param shard_mapping: 分表映射 {table_name: [keys]} :param keys: 主键字段名列表 :param fields: 要查询的字段列表 :param exclude_fields: 要排除的字段列表 :return: (查询结果列表, 查询到的表列表, 查询失败的表列表) """ all_results = [] queried_tables = [] error_tables = [] logger.info(f"开始执行分表查询，涉及 {len(shard_mapping)} 张分表") total_start_time = time.time() for table_name, table_keys in shard_mapping.items(): try: logger.info(f"查询分表 {table_name}，包含 {len(table_keys)} 个key: {table_keys}") # 为每个分表执行查询 table_results = execute_query(session, table_name, keys, fields, table_keys, exclude_fields) all_results.extend(table_results) queried_tables.append(table_name) logger.info(f"分表 {table_name} 查询成功，返回 {len(table_results)} 条记录") except Exception as e: logger.error(f"分表 {table_name} 查询失败: {e}") error_tables.append(table_name) total_execution_time = time.time() - total_start_time logger.info(f"分表查询总计完成: 执行时间={total_execution_time:.3f}秒, 成功表数={len(queried_tables)}, 失败表数={len(error_tables)}, 总记录数={len(all_results)}") return all_results, queried_tables, error_tables def execute_mixed_query(pro_session, test_session, pro_config, test_config, keys, fields_to_compare, values, exclude_fields, sharding_config): """ 执行混合查询（生产环境分表，测试环境可能单表或分表） """ results = { 'pro_data': [], 'test_data': [], 'sharding_info': { 'calculation_stats': {} } } # 处理生产环境查询 if sharding_config.get('use_sharding_for_pro', False): # 获取生产环境分表配置参数，优先使用专用参数，否则使用通用参数 pro_interval = sharding_config.get('pro_interval_seconds') or sharding_config.get('interval_seconds', 604800) pro_table_count = sharding_config.get('pro_table_count') or sharding_config.get('table_count', 14) # 记录生产环境分表配置信息 logger.info(f"=== 生产环境分表配置 ===") logger.info(f"启用分表查询: True") logger.info(f"时间间隔: {pro_interval}秒 ({pro_interval//86400}天)") logger.info(f"分表数量: {pro_table_count}张") logger.info(f"基础表名: {pro_config['table']}") pro_calculator = ShardingCalculator( interval_seconds=pro_interval, table_count=pro_table_count ) pro_shard_mapping, pro_failed_keys, pro_calc_stats = pro_calculator.get_all_shard_tables_for_keys( pro_config['table'], values ) logger.info(f"生产环境分表映射结果: 涉及{len(pro_shard_mapping)}张分表, 失败Key数量: {len(pro_failed_keys)}") pro_data, pro_queried_tables, pro_error_tables = execute_sharding_query( pro_session, pro_shard_mapping, keys, fields_to_compare, exclude_fields ) results['pro_data'] = pro_data results['sharding_info']['pro_shards'] = { 'enabled': True, 'interval_seconds': sharding_config.get('pro_interval_seconds', 604800), 'table_count': sharding_config.get('pro_table_count', 14), 'queried_tables': pro_queried_tables, 'error_tables': pro_error_tables, 'failed_keys': pro_failed_keys } results['sharding_info']['calculation_stats'].update(pro_calc_stats) else: # 生产环境单表查询 logger.info(f"=== 生产环境单表配置 ===") logger.info(f"启用分表查询: False") logger.info(f"表名: {pro_config['table']}") pro_data = execute_query(pro_session, pro_config['table'], keys, fields_to_compare, values, exclude_fields) results['pro_data'] = pro_data results['sharding_info']['pro_shards'] = { 'enabled': False, 'queried_tables': [pro_config['table']] } # 处理测试环境查询 if sharding_config.get('use_sharding_for_test', False): # 获取测试环境分表配置参数，优先使用专用参数，否则使用通用参数 test_interval = sharding_config.get('test_interval_seconds') or sharding_config.get('interval_seconds', 604800) test_table_count = sharding_config.get('test_table_count') or sharding_config.get('table_count', 14) # 记录测试环境分表配置信息 logger.info(f"=== 测试环境分表配置 ===") logger.info(f"启用分表查询: True") logger.info(f"时间间隔: {test_interval}秒 ({test_interval//86400}天)") logger.info(f"分表数量: {test_table_count}张") logger.info(f"基础表名: {test_config['table']}") test_calculator = ShardingCalculator( interval_seconds=test_interval, table_count=test_table_count ) test_shard_mapping, test_failed_keys, test_calc_stats = test_calculator.get_all_shard_tables_for_keys( test_config['table'], values ) logger.info(f"测试环境分表映射结果: 涉及{len(test_shard_mapping)}张分表, 失败Key数量: {len(test_failed_keys)}") test_data, test_queried_tables, test_error_tables = execute_sharding_query( test_session, test_shard_mapping, keys, fields_to_compare, exclude_fields ) results['test_data'] = test_data results['sharding_info']['test_shards'] = { 'enabled': True, 'interval_seconds': test_interval, 'table_count': test_table_count, 'queried_tables': test_queried_tables, 'error_tables': test_error_tables, 'failed_keys': test_failed_keys } # 合并计算统计信息 if not results['sharding_info']['calculation_stats']: results['sharding_info']['calculation_stats'] = test_calc_stats else: # 测试环境单表查询 logger.info(f"=== 测试环境单表配置 ===") logger.info(f"启用分表查询: False") logger.info(f"表名: {test_config['table']}") test_data = execute_query(test_session, test_config['table'], keys, fields_to_compare, values, exclude_fields) results['test_data'] = test_data results['sharding_info']['test_shards'] = { 'enabled': False, 'queried_tables': [test_config['table']] } return results