DataOps-platform/datafactory/scripts/DF_DO202601130001.py

"""
数据流程脚本：DF_DO202601130001
仓库库存汇总表 数据流程

功能：
- 从产品库存表(test_product_inventory)中读取数据
- 按仓库(warehouse)汇总库存数量(current_stock)
- 输出仓库编号和总库存数量到目标表(warehouse_inventory_summary)
- 更新模式：Full Refresh (全量更新)

任务信息：
- DataFlow ID: 2220
- DataFlow Name: 仓库库存汇总表_数据流程
- Order ID: 17
- Order No: DO202601130001

作者：cursor (自动生成)
创建时间：2026-01-13
"""

from __future__ import annotations

import argparse
import logging
import os
import sys
from datetime import datetime
from typing import Any

# 添加项目根目录到路径
sys.path.insert(
    0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
)

try:
    from sqlalchemy import create_engine, text
    from sqlalchemy.orm import sessionmaker
except ImportError:
    print("错误：请安装 sqlalchemy 库")
    sys.exit(1)

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)


class WarehouseInventorySummaryFlow:
    """仓库库存汇总表数据流程处理器"""

    # 配置常量
    SOURCE_TABLE = "test_product_inventory"
    TARGET_TABLE = "warehouse_inventory_summary"
    SOURCE_SCHEMA = "public"
    TARGET_SCHEMA = "public"
    UPDATE_MODE = "full"  # full = 全量更新

    def __init__(self, db_uri: str | None = None):
        """
        初始化数据流程处理器

        Args:
            db_uri: 数据库连接 URI，如果不提供则从配置中获取
        """
        self.db_uri = db_uri or self._get_db_uri()
        self.engine = None
        self.session = None
        self.processed_count = 0
        self.error_count = 0

    def _get_db_uri(self) -> str:
        """获取数据库连接 URI"""
        # 优先从环境变量获取
        db_uri = os.environ.get("DATABASE_URL")
        if db_uri:
            return db_uri

        # 尝试从 Flask 配置获取
        try:
            from app.config.config import config, get_environment

            env = get_environment()
            cfg = config.get(env, config["default"])
            return cfg.SQLALCHEMY_DATABASE_URI
        except ImportError:
            pass

        # 默认使用开发环境配置
        return "postgresql://postgres:postgres@localhost:5432/dataops"

    def connect(self) -> bool:
        """
        连接数据库

        Returns:
            连接是否成功
        """
        try:
            self.engine = create_engine(self.db_uri)
            Session = sessionmaker(bind=self.engine)
            self.session = Session()

            # 测试连接
            with self.engine.connect() as conn:
                conn.execute(text("SELECT 1"))

            # 隐藏密码显示连接信息
            safe_uri = self.db_uri.split("@")[-1] if "@" in self.db_uri else self.db_uri
            logger.info(f"成功连接数据库: {safe_uri}")
            return True

        except Exception as e:
            logger.error(f"连接数据库失败: {str(e)}")
            return False

    def ensure_target_table(self) -> bool:
        """
        确保目标表存在，如果不存在则创建

        Returns:
            操作是否成功
        """
        try:
            if not self.session:
                logger.error("数据库会话未初始化")
                return False

            # 检查目标表是否存在
            check_sql = text("""
                SELECT EXISTS (
                    SELECT FROM information_schema.tables
                    WHERE table_schema = :schema
                    AND table_name = :table_name
                )
            """)

            result = self.session.execute(
                check_sql,
                {"schema": self.TARGET_SCHEMA, "table_name": self.TARGET_TABLE},
            )
            exists = result.scalar()

            if exists:
                logger.info(f"目标表 {self.TARGET_SCHEMA}.{self.TARGET_TABLE} 已存在")
                return True

            # 创建目标表
            create_sql = text(f"""
                CREATE TABLE {self.TARGET_SCHEMA}.{self.TARGET_TABLE} (
                    id SERIAL PRIMARY KEY,
                    warehouse VARCHAR(100) NOT NULL COMMENT '仓库编号',
                    total_stock BIGINT NOT NULL DEFAULT 0 COMMENT '总库存数量',
                    create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP COMMENT '数据创建时间',
                    update_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP COMMENT '数据更新时间'
                );
                COMMENT ON TABLE {self.TARGET_SCHEMA}.{self.TARGET_TABLE}
                    IS '仓库库存汇总表';
                COMMENT ON COLUMN {self.TARGET_SCHEMA}.{self.TARGET_TABLE}.warehouse
                    IS '仓库编号';
                COMMENT ON COLUMN {self.TARGET_SCHEMA}.{self.TARGET_TABLE}.total_stock
                    IS '总库存数量';
                COMMENT ON COLUMN {self.TARGET_SCHEMA}.{self.TARGET_TABLE}.create_time
                    IS '数据创建时间';
                COMMENT ON COLUMN {self.TARGET_SCHEMA}.{self.TARGET_TABLE}.update_time
                    IS '数据更新时间';
            """)

            self.session.execute(create_sql)
            self.session.commit()

            logger.info(f"成功创建目标表: {self.TARGET_SCHEMA}.{self.TARGET_TABLE}")
            return True

        except Exception as e:
            if self.session:
                self.session.rollback()
            logger.error(f"创建目标表失败: {str(e)}")
            # 尝试使用简化的 DDL（PostgreSQL 不支持 COMMENT 在列定义中）
            return self._create_table_simple()

    def _create_table_simple(self) -> bool:
        """使用简化的 DDL 创建目标表（PostgreSQL 兼容）"""
        try:
            if not self.session:
                return False

            # PostgreSQL 简化的建表语句
            create_sql = text(f"""
                CREATE TABLE IF NOT EXISTS {self.TARGET_SCHEMA}.{self.TARGET_TABLE} (
                    id SERIAL PRIMARY KEY,
                    warehouse VARCHAR(100) NOT NULL,
                    total_stock BIGINT NOT NULL DEFAULT 0,
                    create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    update_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)

            self.session.execute(create_sql)

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

            # 添加列注释
            comments = [
                ("warehouse", "仓库编号"),
                ("total_stock", "总库存数量"),
                ("create_time", "数据创建时间"),
                ("update_time", "数据更新时间"),
            ]
            for col_name, col_comment in comments:
                comment_col_sql = text(f"""
                    COMMENT ON COLUMN {self.TARGET_SCHEMA}.{self.TARGET_TABLE}.{col_name}
                    IS '{col_comment}'
                """)
                self.session.execute(comment_col_sql)

            self.session.commit()
            logger.info(
                f"成功创建目标表(简化模式): {self.TARGET_SCHEMA}.{self.TARGET_TABLE}"
            )
            return True

        except Exception as e:
            if self.session:
                self.session.rollback()
            logger.error(f"创建目标表(简化模式)失败: {str(e)}")
            return False

    def extract_and_transform(self) -> list[dict[str, Any]]:
        """
        从源表提取数据并进行转换（按仓库汇总）

        Returns:
            转换后的数据列表
        """
        try:
            if not self.session:
                logger.error("数据库会话未初始化")
                return []

            # 执行汇总查询
            # 1. 从产品库存表中提取字段：仓库编号、产品编号、库存数量
            # 2. 对库存数量进行按仓库编号进行求和计算
            # 3. 无特殊过滤条件
            # 4. 最终输出数据格式包含字段：仓库编号、总库存数量
            query_sql = text(f"""
                SELECT
                    warehouse,
                    SUM(current_stock) as total_stock
                FROM {self.SOURCE_SCHEMA}.{self.SOURCE_TABLE}
                GROUP BY warehouse
                ORDER BY warehouse
            """)

            result = self.session.execute(query_sql)
            rows = result.fetchall()

            data_list = []
            for row in rows:
                data_list.append(
                    {
                        "warehouse": row.warehouse,
                        "total_stock": int(row.total_stock) if row.total_stock else 0,
                    }
                )

            logger.info(f"从源表提取并汇总了 {len(data_list)} 条仓库库存记录")
            return data_list

        except Exception as e:
            logger.error(f"提取和转换数据失败: {str(e)}")
            return []

    def load_to_target(self, data_list: list[dict[str, Any]]) -> bool:
        """
        将转换后的数据加载到目标表

        Args:
            data_list: 转换后的数据列表

        Returns:
            加载是否成功
        """
        try:
            if not data_list:
                logger.warning("没有数据需要加载")
                return True

            if not self.session:
                logger.error("数据库会话未初始化")
                return False

            # 全量更新模式：先清空目标表
            if self.UPDATE_MODE == "full":
                delete_sql = text(
                    f"DELETE FROM {self.TARGET_SCHEMA}.{self.TARGET_TABLE}"
                )
                self.session.execute(delete_sql)
                logger.info(f"目标表 {self.TARGET_TABLE} 已清空（全量更新模式）")

            # 插入新数据
            current_time = datetime.now()
            insert_sql = text(f"""
                INSERT INTO {self.TARGET_SCHEMA}.{self.TARGET_TABLE}
                    (warehouse, total_stock, create_time, update_time)
                VALUES
                    (:warehouse, :total_stock, :create_time, :update_time)
            """)

            for data in data_list:
                try:
                    self.session.execute(
                        insert_sql,
                        {
                            "warehouse": data["warehouse"],
                            "total_stock": data["total_stock"],
                            "create_time": current_time,
                            "update_time": current_time,
                        },
                    )
                    self.processed_count += 1
                except Exception as e:
                    self.error_count += 1
                    logger.error(f"插入数据失败: {str(e)}, 数据: {data}")

            self.session.commit()
            logger.info(
                f"数据加载完成: 成功 {self.processed_count} 条, 失败 {self.error_count} 条"
            )
            return True

        except Exception as e:
            if self.session:
                self.session.rollback()
            logger.error(f"加载数据到目标表失败: {str(e)}")
            return False

    def close(self) -> None:
        """关闭数据库连接"""
        if self.session:
            try:
                self.session.close()
                logger.info("数据库会话已关闭")
            except Exception as e:
                logger.error(f"关闭数据库会话失败: {str(e)}")

        if self.engine:
            try:
                self.engine.dispose()
                logger.info("数据库引擎已释放")
            except Exception as e:
                logger.error(f"释放数据库引擎失败: {str(e)}")

    def run(self) -> dict[str, Any]:
        """
        执行完整的 ETL 流程

        Returns:
            执行结果字典
        """
        result = {
            "success": False,
            "processed_count": 0,
            "error_count": 0,
            "update_mode": self.UPDATE_MODE,
            "source_table": f"{self.SOURCE_SCHEMA}.{self.SOURCE_TABLE}",
            "target_table": f"{self.TARGET_SCHEMA}.{self.TARGET_TABLE}",
            "message": "",
        }

        try:
            logger.info("=" * 60)
            logger.info("开始执行数据流程: DF_DO202601130001")
            logger.info(f"源表: {self.SOURCE_SCHEMA}.{self.SOURCE_TABLE}")
            logger.info(f"目标表: {self.TARGET_SCHEMA}.{self.TARGET_TABLE}")
            logger.info(f"更新模式: {self.UPDATE_MODE}")
            logger.info("=" * 60)

            # 1. 连接数据库
            if not self.connect():
                result["message"] = "连接数据库失败"
                return result

            # 2. 确保目标表存在
            if not self.ensure_target_table():
                result["message"] = "创建目标表失败"
                return result

            # 3. 提取和转换数据
            data_list = self.extract_and_transform()

            if not data_list:
                result["message"] = "未提取到数据"
                result["success"] = True  # 没有数据不算失败
                return result

            # 4. 加载到目标表
            if self.load_to_target(data_list):
                result["success"] = True
                result["processed_count"] = self.processed_count
                result["error_count"] = self.error_count
                result["message"] = (
                    f"数据流程执行成功: "
                    f"处理 {self.processed_count} 条, 失败 {self.error_count} 条"
                )
            else:
                result["message"] = "加载数据到目标表失败"

        except Exception as e:
            logger.error(f"数据流程执行异常: {str(e)}")
            result["message"] = f"数据流程执行异常: {str(e)}"

        finally:
            self.close()

        logger.info("=" * 60)
        logger.info(f"执行结果: {result['message']}")
        logger.info("=" * 60)

        return result


def main():
    """主函数"""
    parser = argparse.ArgumentParser(
        description="DF_DO202601130001 - 仓库库存汇总表数据流程"
    )
    parser.add_argument(
        "--db-uri",
        type=str,
        default=None,
        help="数据库连接 URI (可选，默认从配置获取)",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="仅测试连接和查询，不执行写入",
    )

    args = parser.parse_args()

    # 创建并执行数据流程
    flow = WarehouseInventorySummaryFlow(db_uri=args.db_uri)

    if args.dry_run:
        logger.info("Dry-run 模式: 仅测试连接和查询")
        if flow.connect():
            data_list = flow.extract_and_transform()
            logger.info(f"预览数据 ({len(data_list)} 条):")
            for data in data_list:
                logger.info(f"  {data}")
            flow.close()
            print("\nDry-run 完成，未执行写入操作")
            sys.exit(0)
        else:
            print("\n连接失败")
            sys.exit(1)

    result = flow.run()

    # 输出结果
    print("\n" + "=" * 60)
    print(f"数据流程执行结果: {'成功' if result['success'] else '失败'}")
    print(f"消息: {result['message']}")
    print(f"处理记录数: {result['processed_count']}")
    print(f"失败记录数: {result['error_count']}")
    print(f"更新模式: {result['update_mode']}")
    print(f"源表: {result['source_table']}")
    print(f"目标表: {result['target_table']}")
    print("=" * 60)

    # 设置退出代码
    sys.exit(0 if result["success"] else 1)


if __name__ == "__main__":
    main()