DataOps-platform/datafactory/scripts/task_44_DF_DO202601210001.py

#!/usr/bin/env python
"""
数据流任务脚本 - 仓库库存汇总表_数据流程

任务ID: 44
任务名称: DF_DO202601210001
创建时间: 2026-01-21
更新模式: Append (追加模式)

关联信息:
- Order ID: 26
- Order No: DO202601210001
- DataFlow ID: 2291
- DataFlow Name: 仓库库存汇总表_数据流程
- Product ID: 23

描述: 从源表 test_product_inventory 提取仓库名称和库存数量，
按仓库名称分组求和，输出到 warehouse_inventory_summary 表
"""

from __future__ import annotations

import os
import sys
from datetime import datetime
from typing import Any

import pandas as pd
import psycopg2
from loguru import logger

# 添加项目根目录到Python路径
PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))
sys.path.insert(0, PROJECT_ROOT)

# 任务配置
TASK_ID = 44
TASK_NAME = "DF_DO202601210001"
UPDATE_MODE = "append"

# 源数据库配置（与目标相同）
SOURCE_CONFIG = {
    "host": "192.168.3.143",
    "port": 5432,  # PostgreSQL 默认端口
    "database": "dataops",
    "user": "postgres",
    "password": "dataOps",
}

# 目标数据库配置
TARGET_CONFIG = {
    "host": "192.168.3.143",
    "port": 5432,  # PostgreSQL 默认端口
    "database": "dataops",
    "user": "postgres",
    "password": "dataOps",
}

# 源表配置
SOURCE_SCHEMA = "dags"
SOURCE_TABLE = "test_product_inventory"

# 目标表配置
TARGET_SCHEMA = "dags"
TARGET_TABLE = "warehouse_inventory_summary"


def get_source_connection() -> psycopg2.extensions.connection:
    """
    获取源数据库连接

    Returns:
        psycopg2 连接对象
    """
    conn = psycopg2.connect(
        host=SOURCE_CONFIG["host"],
        port=SOURCE_CONFIG["port"],
        database=SOURCE_CONFIG["database"],
        user=SOURCE_CONFIG["user"],
        password=SOURCE_CONFIG["password"],
        options=f"-c search_path={SOURCE_SCHEMA},public",
    )
    logger.info("源数据库连接成功")
    return conn


def get_target_connection() -> psycopg2.extensions.connection:
    """
    获取目标数据库连接

    Returns:
        psycopg2 连接对象
    """
    conn = psycopg2.connect(
        host=TARGET_CONFIG["host"],
        port=TARGET_CONFIG["port"],
        database=TARGET_CONFIG["database"],
        user=TARGET_CONFIG["user"],
        password=TARGET_CONFIG["password"],
        options=f"-c search_path={TARGET_SCHEMA},public",
    )
    logger.info("目标数据库连接成功")
    return conn


def ensure_target_table_exists(conn: psycopg2.extensions.connection) -> None:
    """
    确保目标表存在，如果不存在则创建

    Args:
        conn: 目标数据库连接
    """
    cursor = conn.cursor()

    try:
        # 检查表是否存在
        cursor.execute(
            """
            SELECT EXISTS(
                SELECT 1 FROM information_schema.tables
                WHERE table_schema = %s
                AND table_name = %s
            )
        """,
            (TARGET_SCHEMA, TARGET_TABLE),
        )
        result = cursor.fetchone()
        exists = result[0] if result else False

        if not exists:
            logger.info(f"目标表不存在，正在创建 {TARGET_SCHEMA}.{TARGET_TABLE}...")

            # 根据任务描述中的 DDL 创建表
            create_table_sql = f"""
            CREATE TABLE IF NOT EXISTS {TARGET_SCHEMA}.{TARGET_TABLE} (
                warehouse_name VARCHAR(255),
                total_inventory INTEGER,
                create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            );
            """
            cursor.execute(create_table_sql)

            # 添加表注释
            cursor.execute(
                f"COMMENT ON TABLE {TARGET_SCHEMA}.{TARGET_TABLE} IS '仓库库存汇总表'"
            )

            # 添加列注释
            column_comments = {
                "warehouse_name": "仓库名称",
                "total_inventory": "库存数量",
                "create_time": "数据创建时间",
            }
            for col, comment in column_comments.items():
                try:
                    cursor.execute(
                        f"COMMENT ON COLUMN {TARGET_SCHEMA}.{TARGET_TABLE}.{col} IS %s",
                        (comment,),
                    )
                except Exception as e:
                    logger.warning(f"添加列注释失败 {col}: {e}")

            conn.commit()
            logger.info(f"目标表 {TARGET_SCHEMA}.{TARGET_TABLE} 创建成功")
        else:
            logger.info(f"目标表 {TARGET_SCHEMA}.{TARGET_TABLE} 已存在")

    except Exception as e:
        conn.rollback()
        logger.error(f"创建目标表失败: {e}")
        raise
    finally:
        cursor.close()


def extract_source_data(conn: psycopg2.extensions.connection) -> pd.DataFrame:
    """
    从源表提取数据

    根据任务描述：
    1. 从源数据'产品库存表'中提取'仓库名称'字段；
    2. 对'产品库存表'中的'库存数量'字段进行求和计算；
    3. 按'仓库名称'进行分组；
    4. 最终输出数据格式包含'仓库名称'和对应的'库存数量'两个字段。

    Args:
        conn: 源数据库连接

    Returns:
        包含汇总数据的DataFrame
    """
    query = f"""
    SELECT
        warehouse AS warehouse_name,
        SUM(current_stock) AS total_inventory
    FROM {SOURCE_SCHEMA}.{SOURCE_TABLE}
    WHERE warehouse IS NOT NULL
    GROUP BY warehouse
    ORDER BY warehouse
    """

    logger.info("正在从源表提取并汇总数据...")

    try:
        df = pd.read_sql(query, conn)
        logger.info(f"成功提取 {len(df)} 条汇总记录")
        return df
    except Exception as e:
        logger.error(f"提取源数据失败: {e}")
        raise


def transform_data(df: pd.DataFrame) -> pd.DataFrame:
    """
    数据转换处理

    Args:
        df: 源数据DataFrame

    Returns:
        转换后的DataFrame
    """
    logger.info("正在执行数据转换...")

    # 数据已在SQL中完成汇总，此处仅做数据清洗
    # 确保字段名称与目标表一致
    df = df.rename(
        columns={
            "warehouse_name": "warehouse_name",
            "total_inventory": "total_inventory",
        }
    )

    # 处理空值：将 None/NaN 的库存数量设为 0
    df["total_inventory"] = df["total_inventory"].fillna(0).astype(int)

    logger.info(f"数据转换完成，共 {len(df)} 条记录")
    return df


def load_to_target(
    df: pd.DataFrame,
    conn: psycopg2.extensions.connection,
    update_mode: str = "append",
    batch_size: int = 1000,
) -> int:
    """
    将数据加载到目标表

    Args:
        df: 要加载的DataFrame
        conn: 目标数据库连接
        update_mode: 更新模式（append 或 full）
        batch_size: 批量插入大小

    Returns:
        插入的记录数
    """
    if df.empty:
        logger.warning("没有数据需要加载")
        return 0

    logger.info(f"正在将 {len(df)} 条记录加载到目标表...")

    target_table = f"{TARGET_SCHEMA}.{TARGET_TABLE}"
    cursor = conn.cursor()
    inserted_count = 0

    try:
        # 全量更新模式：先清空目标表
        if update_mode.lower() == "full":
            logger.info("全量更新模式：清空目标表...")
            cursor.execute(f"TRUNCATE TABLE {target_table}")
            logger.info("目标表已清空")

        # 目标表结构准备插入的列（不包含 create_time，由数据库自动设置）
        columns = ["warehouse_name", "total_inventory"]

        # 构建插入SQL
        placeholders = ", ".join(["%s"] * len(columns))
        column_names = ", ".join(columns)
        insert_sql = (
            f"INSERT INTO {target_table} ({column_names}) VALUES ({placeholders})"
        )

        # 批量插入
        for i in range(0, len(df), batch_size):
            batch_df = df.iloc[i : i + batch_size]
            records = []
            for _, row in batch_df.iterrows():
                record = tuple(
                    None if pd.isna(row.get(col)) else row.get(col) for col in columns
                )
                records.append(record)

            cursor.executemany(insert_sql, records)
            inserted_count += len(records)
            logger.debug(f"已插入 {inserted_count}/{len(df)} 条记录")

        conn.commit()
        logger.info(f"成功加载 {inserted_count} 条记录到 {target_table}")
        return inserted_count

    except Exception as e:
        conn.rollback()
        logger.error(f"数据加载失败: {e}")
        raise
    finally:
        cursor.close()


def main() -> dict[str, Any]:
    """
    主函数：执行ETL流程

    Returns:
        执行结果字典
    """
    result = {
        "task_id": TASK_ID,
        "task_name": TASK_NAME,
        "status": "failed",
        "records_extracted": 0,
        "records_loaded": 0,
        "error_message": None,
        "execution_time": None,
    }

    start_time = datetime.now()
    source_conn = None
    target_conn = None

    try:
        logger.info("=" * 60)
        logger.info(f"任务开始: {TASK_NAME}")
        logger.info("=" * 60)

        # 步骤1: 建立数据库连接
        logger.info("[Step 1/5] 建立数据库连接...")
        source_conn = get_source_connection()
        target_conn = get_target_connection()

        # 步骤2: 确保目标表存在（重要：必须在数据加载前执行）
        logger.info("[Step 2/5] 检查/创建目标表...")
        ensure_target_table_exists(target_conn)

        # 步骤3: 从源表提取数据
        logger.info("[Step 3/5] 提取并汇总源数据...")
        df = extract_source_data(source_conn)
        result["records_extracted"] = len(df)

        # 步骤4: 数据转换
        logger.info("[Step 4/5] 数据转换...")
        df_transformed = transform_data(df)

        # 步骤5: 加载到目标表
        logger.info("[Step 5/5] 加载数据到目标表...")
        records_loaded = load_to_target(
            df_transformed, target_conn, update_mode=UPDATE_MODE
        )
        result["records_loaded"] = records_loaded

        result["status"] = "success"
        logger.info("=" * 60)
        logger.info(
            f"任务完成! 提取: {result['records_extracted']}, 加载: {result['records_loaded']}"
        )
        logger.info("=" * 60)

    except Exception as e:
        result["status"] = "failed"
        result["error_message"] = str(e)
        logger.error(f"任务执行失败: {e}")
        raise

    finally:
        # 关闭数据库连接
        if source_conn:
            source_conn.close()
            logger.debug("源数据库连接已关闭")
        if target_conn:
            target_conn.close()
            logger.debug("目标数据库连接已关闭")

        result["execution_time"] = str(datetime.now() - start_time)

    return result


if __name__ == "__main__":
    # 配置日志
    # 重要：日志输出到 stdout 而非 stderr，以便 n8n 工作流正确解析输出
    logger.remove()
    logger.add(
        sys.stdout,
        level="INFO",
        format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>",
    )
    logger.add(
        os.path.join(PROJECT_ROOT, "logs", f"task_{TASK_ID}.log"),
        level="DEBUG",
        rotation="10 MB",
        retention="7 days",
        encoding="utf-8",
    )

    try:
        result = main()
        if result["status"] == "success":
            sys.exit(0)
        else:
            sys.exit(1)
    except Exception as e:
        logger.exception(f"脚本执行异常: {e}")
        sys.exit(1)