""" 数据资源导入工具 功能：从远程数据源读取数据，按照指定的更新模式写入到目标数据资源表中 支持： - 灵活的数据源配置（PostgreSQL/MySQL等） - 灵活的目标表配置 - 两种更新模式：append（追加）/ full（全量更新） 作者：cursor 创建时间：2025-11-28 更新时间：2025-11-28 """ import argparse import json import logging import os import sys from typing import Any, Dict, List, Optional import psycopg2 from sqlalchemy import create_engine, inspect, text from sqlalchemy.engine import Engine from sqlalchemy.orm import Session, sessionmaker # 添加项目根目录到路径 sys.path.insert( 0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))) ) try: from app.config.config import config, get_environment # type: ignore # 获取当前环境的配置类 _current_env = get_environment() Config = config.get(_current_env, config["default"]) except ImportError: # 如果无法导入，使用环境变量 class Config: # type: ignore SQLALCHEMY_DATABASE_URI = os.environ.get( "DATABASE_URI", "postgresql://user:password@localhost:5432/database" ) try: import pymysql # type: ignore MYSQL_AVAILABLE = True except ImportError: MYSQL_AVAILABLE = False pymysql = None # type: ignore # 配置日志 logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" ) logger = logging.getLogger(__name__) class ResourceDataImporter: """数据资源导入器""" # 目标表所在的 schema TARGET_SCHEMA = "dags" def __init__( self, source_config: Dict[str, Any], target_table_name: str, update_mode: str = "append", ): """ 初始化导入器 Args: source_config: 源数据库配置 { 'type': 'postgresql', # 或 'mysql' 'host': '10.52.31.104', 'port': 5432, 'database': 'source_db', 'username': 'user', 'password': 'password', 'table_name': 'TB_JC_KSDZB' # 源表名 } target_table_name: 目标表名（数据资源的英文名） update_mode: 更新模式，'append'（追加）或 'full'（全量更新） """ self.source_config = source_config self.target_table_name = target_table_name self.update_mode = update_mode.lower() self.source_connection: Optional[Any] = None self.target_engine: Optional[Engine] = None self.target_session: Optional[Session] = None self.imported_count = 0 self.updated_count = 0 self.error_count = 0 # 验证更新模式 if self.update_mode not in ["append", "full"]: raise ValueError( f"不支持的更新模式: {update_mode}，仅支持 'append' 或 'full'" ) logger.info( f"初始化数据导入器: 目标表={self.TARGET_SCHEMA}.{target_table_name}, 更新模式={update_mode}" ) def connect_target_database(self) -> bool: """ 连接目标数据库（从 config.py 获取配置） Returns: 连接是否成功 """ try: # 从 Config 获取 PostgreSQL 配置 db_uri = Config.SQLALCHEMY_DATABASE_URI if not db_uri: logger.error("未找到目标数据库配置（SQLALCHEMY_DATABASE_URI）") return False # 创建目标数据库引擎 self.target_engine = create_engine(db_uri) Session = sessionmaker(bind=self.target_engine) self.target_session = Session() # 测试连接 self.target_engine.connect() logger.info(f"成功连接目标数据库: {db_uri.split('@')[-1]}") # 隐藏密码 return True except Exception as e: logger.error(f"连接目标数据库失败: {str(e)}") return False def connect_source_database(self) -> bool: """ 连接源数据库 Returns: 连接是否成功 """ try: db_type = self.source_config["type"].lower() if db_type == "postgresql": self.source_connection = psycopg2.connect( host=self.source_config["host"], port=self.source_config["port"], database=self.source_config["database"], user=self.source_config["username"], password=self.source_config["password"], ) logger.info( f"成功连接源数据库（PostgreSQL）: {self.source_config['host']}:{self.source_config['port']}/{self.source_config['database']}" ) return True elif db_type == "mysql": if not MYSQL_AVAILABLE or pymysql is None: logger.error("pymysql未安装，无法连接MySQL数据库") return False self.source_connection = pymysql.connect( host=self.source_config["host"], port=self.source_config["port"], database=self.source_config["database"], user=self.source_config["username"], password=self.source_config["password"], ) logger.info( f"成功连接源数据库（MySQL）: {self.source_config['host']}:{self.source_config['port']}/{self.source_config['database']}" ) return True else: logger.error(f"不支持的数据库类型: {db_type}") return False except Exception as e: logger.error(f"连接源数据库失败: {str(e)}") return False def get_full_table_name(self) -> str: """ 获取带 schema 的完整表名 Returns: 完整表名 (schema.table_name) """ return f"{self.TARGET_SCHEMA}.{self.target_table_name}" def get_target_table_columns(self) -> List[str]: """ 获取目标表的列名 Returns: 列名列表 """ try: if not self.target_engine: logger.error("目标数据库引擎未初始化") return [] inspector = inspect(self.target_engine) # 指定 schema 来获取表的列名 columns = inspector.get_columns( self.target_table_name, schema=self.TARGET_SCHEMA ) column_names = [ col["name"] for col in columns if col["name"] != "create_time" ] logger.info(f"目标表 {self.get_full_table_name()} 的列: {column_names}") return column_names except Exception as e: logger.error(f"获取目标表列名失败: {str(e)}") return [] def extract_source_data(self, limit: Optional[int] = None) -> List[Dict[str, Any]]: """ 从源数据库提取数据 Args: limit: 限制提取的数据行数（None 表示不限制） Returns: 数据行列表 """ try: if not self.source_connection: logger.error("源数据库连接未建立") return [] cursor = self.source_connection.cursor() source_table = self.source_config.get("table_name") if not source_table: logger.error("源表名未指定") return [] # 构建查询语句 query = f"SELECT * FROM {source_table}" # 添加过滤条件（如果有） where_clause = self.source_config.get("where_clause", "") if where_clause: query += f" WHERE {where_clause}" # 添加排序（如果有） order_by = self.source_config.get("order_by", "") if order_by: query += f" ORDER BY {order_by}" # 添加限制 if limit: query += f" LIMIT {limit}" logger.info(f"执行查询: {query}") cursor.execute(query) # 获取列名 columns = [desc[0] for desc in cursor.description] # 提取数据 rows = [] for row in cursor.fetchall(): row_dict = dict(zip(columns, row)) rows.append(row_dict) cursor.close() logger.info(f"从源表 {source_table} 提取了 {len(rows)} 条数据") return rows except Exception as e: logger.error(f"提取源数据失败: {str(e)}") return [] def clear_target_table(self) -> bool: """ 清空目标表（用于全量更新模式） Returns: 清空是否成功 """ try: if not self.target_session: logger.error("目标数据库会话未初始化") return False full_table_name = self.get_full_table_name() delete_sql = text(f"DELETE FROM {full_table_name}") self.target_session.execute(delete_sql) self.target_session.commit() logger.info(f"目标表 {full_table_name} 已清空") return True except Exception as e: if self.target_session: self.target_session.rollback() logger.error(f"清空目标表失败: {str(e)}") return False def map_source_to_target_columns( self, source_row: Dict[str, Any], target_columns: List[str] ) -> Dict[str, Any]: """ 将源数据列映射到目标表列 Args: source_row: 源数据行 target_columns: 目标表列名列表 Returns: 映射后的数据行 """ mapped_row = {} for col in target_columns: # 优先使用精确匹配（不区分大小写） col_lower = col.lower() for source_col, value in source_row.items(): if source_col.lower() == col_lower: mapped_row[col] = value break else: # 如果没有匹配到，设置为 None mapped_row[col] = None return mapped_row def insert_data_to_target(self, data_rows: List[Dict[str, Any]]) -> bool: """ 将数据插入目标表 Args: data_rows: 数据行列表 Returns: 插入是否成功 """ try: if not data_rows: logger.warning("没有数据需要插入") return True if not self.target_session: logger.error("目标数据库会话未初始化") return False # 获取目标表列名 target_columns = self.get_target_table_columns() if not target_columns: logger.error("无法获取目标表列名") return False # 全量更新模式：先清空目标表 if self.update_mode == "full" and not self.clear_target_table(): return False # 构建插入 SQL（使用带 schema 的完整表名） full_table_name = self.get_full_table_name() columns_str = ", ".join(target_columns + ["create_time"]) placeholders = ", ".join( [f":{col}" for col in target_columns] + ["CURRENT_TIMESTAMP"] ) insert_sql = text(f""" INSERT INTO {full_table_name} ({columns_str}) VALUES ({placeholders}) """) # 批量插入 success_count = 0 for source_row in data_rows: try: # 映射列名 mapped_row = self.map_source_to_target_columns( source_row, target_columns ) # 执行插入 self.target_session.execute(insert_sql, mapped_row) success_count += 1 # 每 100 条提交一次 if success_count % 100 == 0: self.target_session.commit() logger.info(f"已插入 {success_count} 条数据...") except Exception as e: self.error_count += 1 logger.error(f"插入数据失败: {str(e)}, 数据: {source_row}") # 最终提交 self.target_session.commit() self.imported_count = success_count logger.info( f"数据插入完成: 成功 {self.imported_count} 条, 失败 {self.error_count} 条" ) return True except Exception as e: if self.target_session: self.target_session.rollback() logger.error(f"批量插入数据失败: {str(e)}") return False def close_connections(self): """关闭所有数据库连接""" # 关闭源数据库连接 if self.source_connection: try: self.source_connection.close() logger.info("源数据库连接已关闭") except Exception as e: logger.error(f"关闭源数据库连接失败: {str(e)}") # 关闭目标数据库连接 if self.target_session: try: self.target_session.close() logger.info("目标数据库会话已关闭") except Exception as e: logger.error(f"关闭目标数据库会话失败: {str(e)}") if self.target_engine: try: self.target_engine.dispose() logger.info("目标数据库引擎已释放") except Exception as e: logger.error(f"释放目标数据库引擎失败: {str(e)}") def run(self, limit: Optional[int] = None) -> Dict[str, Any]: """ 执行导入流程 Args: limit: 限制导入的数据行数（None 表示不限制） Returns: 执行结果 """ result = { "success": False, "imported_count": 0, "error_count": 0, "update_mode": self.update_mode, "message": "", } try: logger.info("=" * 60) logger.info("开始数据导入") logger.info(f"源表: {self.source_config.get('table_name')}") logger.info(f"目标表: {self.get_full_table_name()}") logger.info(f"更新模式: {self.update_mode}") logger.info("=" * 60) # 1. 连接源数据库 if not self.connect_source_database(): result["message"] = "连接源数据库失败" return result # 2. 连接目标数据库 if not self.connect_target_database(): result["message"] = "连接目标数据库失败" return result # 3. 提取源数据 data_rows = self.extract_source_data(limit=limit) if not data_rows: result["message"] = "未提取到数据" result["success"] = True # 没有数据不算失败 return result # 4. 插入数据到目标表 if self.insert_data_to_target(data_rows): result["success"] = True result["imported_count"] = self.imported_count result["error_count"] = self.error_count result["message"] = ( f"导入完成: 成功 {self.imported_count} 条, 失败 {self.error_count} 条" ) else: result["message"] = "插入数据到目标表失败" except Exception as e: logger.error(f"导入过程发生异常: {str(e)}") result["message"] = f"导入失败: {str(e)}" finally: # 5. 关闭连接 self.close_connections() logger.info("=" * 60) logger.info(f"导入结果: {result['message']}") logger.info("=" * 60) return result def import_resource_data( source_config: Dict[str, Any], target_table_name: str, update_mode: str = "append", limit: Optional[int] = None, ) -> Dict[str, Any]: """ 导入数据资源（入口函数） Args: source_config: 源数据库配置 { 'type': 'postgresql', # 或 'mysql' 'host': '10.52.31.104', 'port': 5432, 'database': 'source_db', 'username': 'user', 'password': 'password', 'table_name': 'TB_JC_KSDZB', # 源表名 'where_clause': "TBRQ >= '2025-01-01'", # 可选：WHERE条件 'order_by': 'TBRQ DESC' # 可选：排序 } target_table_name: 目标表名（数据资源的英文名） update_mode: 更新模式，'append'（追加）或 'full'（全量更新） limit: 限制导入的数据行数（None 表示不限制） Returns: 导入结果 """ importer = ResourceDataImporter( source_config=source_config, target_table_name=target_table_name, update_mode=update_mode, ) return importer.run(limit=limit) def parse_args(): """解析命令行参数""" parser = argparse.ArgumentParser(description="数据资源导入工具") parser.add_argument( "--source-config", type=str, required=True, help="源数据库配置（JSON格式字符串或文件路径）", ) parser.add_argument( "--target-table", type=str, required=True, help="目标表名（数据资源的英文名）" ) parser.add_argument( "--update-mode", type=str, choices=["append", "full"], default="append", help="更新模式：append（追加）或 full（全量更新）", ) parser.add_argument("--limit", type=int, default=None, help="限制导入的数据行数") return parser.parse_args() if __name__ == "__main__": # 解析命令行参数 args = parse_args() # 解析源数据库配置 try: # 尝试作为JSON字符串解析 source_config = json.loads(args.source_config) except json.JSONDecodeError: # 尝试作为文件路径读取 try: with open(args.source_config, encoding="utf-8") as f: source_config = json.load(f) except Exception as e: logger.error(f"解析源数据库配置失败: {str(e)}") exit(1) # 执行导入 result = import_resource_data( source_config=source_config, target_table_name=args.target_table, update_mode=args.update_mode, limit=args.limit, ) # 输出结果 print("\n" + "=" * 60) print(f"导入结果: {'成功' if result['success'] else '失败'}") print(f"消息: {result['message']}") print(f"成功: {result['imported_count']} 条") print(f"失败: {result['error_count']} 条") print(f"更新模式: {result['update_mode']}") print("=" * 60) # 设置退出代码 exit(0 if result["success"] else 1)