airflow_scheduler/dags/common.py

# common.py
import psycopg2
from neo4j import GraphDatabase
import logging
import importlib.util
from pathlib import Path
import networkx as nx
import os
from datetime import datetime, timedelta
from config import PG_CONFIG, NEO4J_CONFIG, SCRIPTS_BASE_PATH

# 创建统一的日志记录器
logger = logging.getLogger("airflow.task")

def get_pg_conn():
    """获取PostgreSQL连接"""
    return psycopg2.connect(**PG_CONFIG)

def get_neo4j_driver():
    """获取Neo4j连接驱动"""
    uri = NEO4J_CONFIG['uri']
    auth = (NEO4J_CONFIG['user'], NEO4J_CONFIG['password'])
    return GraphDatabase.driver(uri, auth=auth)

def update_task_start_time(exec_date, target_table, script_name, start_time):
    """更新任务开始时间"""
    conn = get_pg_conn()
    cursor = conn.cursor()
    try:
        cursor.execute("""
            UPDATE airflow_dag_schedule 
            SET exec_start_time = %s
            WHERE exec_date = %s AND target_table = %s AND script_name = %s
        """, (start_time, exec_date, target_table, script_name))
        conn.commit()
    except Exception as e:
        logger.error(f"更新任务开始时间失败: {str(e)}")
        conn.rollback()
    finally:
        cursor.close()
        conn.close()

def update_task_completion(exec_date, target_table, script_name, success, end_time, duration):
    """更新任务完成信息"""
    conn = get_pg_conn()
    cursor = conn.cursor()
    try:
        cursor.execute("""
            UPDATE airflow_dag_schedule 
            SET exec_result = %s, exec_end_time = %s, exec_duration = %s
            WHERE exec_date = %s AND target_table = %s AND script_name = %s
        """, (success, end_time, duration, exec_date, target_table, script_name))
        conn.commit()
    except Exception as e:
        logger.error(f"更新任务完成信息失败: {str(e)}")
        conn.rollback()
    finally:
        cursor.close()
        conn.close()

def execute_with_monitoring(target_table, script_name, script_exec_mode, exec_date, **kwargs):
    """执行脚本并监控执行情况"""

    # 检查script_name是否为空
    if not script_name:
        logger.error(f"表 {target_table} 的script_name为空，无法执行")
        # 记录执行失败
        now = datetime.now()
        update_task_completion(exec_date, target_table, script_name or "", False, now, 0)
        return False
    # 记录执行开始时间
    start_time = datetime.now()
    update_task_start_time(exec_date, target_table, script_name, start_time)
    
    try:
        # 执行实际脚本
        success = execute_script(script_name, target_table, script_exec_mode)
        
        # 记录结束时间和结果
        end_time = datetime.now()
        duration = (end_time - start_time).total_seconds()
        update_task_completion(exec_date, target_table, script_name, success, end_time, duration)
        
        return success
    except Exception as e:
        # 处理异常
        logger.error(f"执行任务出错: {str(e)}")
        end_time = datetime.now()
        duration = (end_time - start_time).total_seconds()
        update_task_completion(exec_date, target_table, script_name, False, end_time, duration)
        raise e

def execute_script(script_name, table_name, execution_mode):
    """执行脚本并返回结果"""
    if not script_name:
        logger.error("未提供脚本名称，无法执行")
        return False
    
    try:
        # 直接使用配置的部署路径
        script_path = Path(SCRIPTS_BASE_PATH) / script_name
        logger.info(f"使用配置的Airflow部署路径: {script_path}")
        
        # 动态导入模块
        spec = importlib.util.spec_from_file_location("dynamic_module", script_path)
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)
        
        # 使用标准入口函数run
        if hasattr(module, "run"):
            logger.info(f"执行脚本 {script_name} 的标准入口函数 run()")
            result = module.run(table_name=table_name, execution_mode=execution_mode)
            logger.info(f"脚本 {script_name} 执行结果: {result}")
            return result
        else:
            logger.warning(f"脚本 {script_name} 未定义标准入口函数 run()，无法执行")
            return False
    except Exception as e:
        logger.error(f"执行脚本 {script_name} 时出错: {str(e)}")
        import traceback
        logger.error(traceback.format_exc())
        return False

def generate_optimized_execution_order(table_names, dependency_dict):
    """
    生成优化的执行顺序，处理循环依赖
    
    参数:
        table_names: 表名列表
        dependency_dict: 依赖关系字典 {表名: [依赖表1, 依赖表2, ...]}
    
    返回:
        list: 优化后的执行顺序列表
    """
    # 创建有向图
    G = nx.DiGraph()
    
    # 添加所有节点
    for table_name in table_names:
        G.add_node(table_name)
    
    # 添加依赖边
    for target, sources in dependency_dict.items():
        for source in sources:
            if source in table_names:  # 确保只考虑目标表集合中的表
                # 从依赖指向目标，表示依赖需要先执行
                G.add_edge(source, target)
    
    # 检测循环依赖
    cycles = list(nx.simple_cycles(G))
    if cycles:
        logger.warning(f"检测到循环依赖，将尝试打破循环: {cycles}")
        # 打破循环依赖（简单策略：移除每个循环中的一条边）
        for cycle in cycles:
            # 移除循环中的最后一条边
            G.remove_edge(cycle[-1], cycle[0])
            logger.info(f"打破循环依赖: 移除 {cycle[-1]} -> {cycle[0]} 的依赖")
    
    # 生成拓扑排序
    try:
        execution_order = list(nx.topological_sort(G))
        return execution_order
    except Exception as e:
        logger.error(f"生成执行顺序失败: {str(e)}")
        # 返回原始列表作为备选
        return table_names

def get_datamodel_dependency_from_neo4j(table_names):
    """
    从Neo4j获取DataModel表间的依赖关系
    
    参数:
        table_names: 表名列表
    
    返回:
        dict: 依赖关系字典 {目标表: [依赖表1, 依赖表2, ...]}
    """
    logger.info(f"开始获取 {len(table_names)} 个表的依赖关系")
    
    # 创建Neo4j连接
    driver = get_neo4j_driver()
    dependency_dict = {name: [] for name in table_names}
    
    try:
        with driver.session() as session:
            # 使用一次性查询获取所有表之间的依赖关系
            query = """
                MATCH (source:DataModel)-[:DERIVED_FROM]->(target:DataModel)
                WHERE source.en_name IN $table_names AND target.en_name IN $table_names
                RETURN source.en_name AS source, target.en_name AS target
            """
            result = session.run(query, table_names=table_names)
            
            # 处理结果
            for record in result:
                source = record.get("source")
                target = record.get("target")
                
                if source and target:
                    # 目标依赖于源
                    if source in dependency_dict:
                        dependency_dict[source].append(target)
                        logger.debug(f"依赖关系: {source} 依赖于 {target}")
    except Exception as e:
        logger.error(f"从Neo4j获取依赖关系时出错: {str(e)}")
    finally:
        driver.close()
    
    # 记录依赖关系
    for table, deps in dependency_dict.items():
        if deps:
            logger.info(f"表 {table} 依赖于: {deps}")
        else:
            logger.info(f"表 {table} 没有依赖")
    
    return dependency_dict

def get_today_date():
    """获取今天的日期，返回YYYY-MM-DD格式字符串"""
    return datetime.now().strftime("%Y-%m-%d")