KCGL/数据爬取/导入基础信息.py

import pandas as pd
import psycopg2
from psycopg2 import sql

# 1. 数据库配置 (根据您的 docker-compose 配置)
DB_CONFIG = {
    'dbname': 'inventory_system',
    'user': 'test',
    'password': '1234',
    'host': 'localhost',  # 脚本在宿主机运行，连接映射出的端口
    'port': '5434'
}

# 2. Excel 文件路径
EXCEL_FILE = '../product.template.xlsx'


def process_excel_to_db():
    try:
        # 读取 Excel 文件
        # dtype=str 确保所有数据读取为字符串，避免数字被自动转换为浮点数（如条码）
        df = pd.read_excel(EXCEL_FILE, dtype=str)

        # 将 NaN (空值) 替换为 None，方便插入数据库时转为 NULL
        df = df.where(pd.notnull(df), None)

        print(f"成功读取 Excel 文件，共 {len(df)} 行数据。")

        # 连接数据库
        conn = psycopg2.connect(**DB_CONFIG)
        cur = conn.cursor()

        success_count = 0

        for index, row in df.iterrows():
            # --- 数据提取与处理逻辑 ---

            # 1. 名称 -> name, common_name
            name = row.get('名称')
            common_name = row.get('名称')  # 题目要求写入到俗名中也用名称

            # 2. 产品类别 -> category
            category = row.get('产品类别')

            # 3. 产品类型 -> material_type
            material_type = row.get('产品类型')

            # 4. 规格型号 -> spec_model (逻辑：条码/内部参考，如果没有内部参考则没有/)
            barcode = row.get('条码')
            internal_ref = row.get('内部参考')

            spec_model = ""
            if barcode and internal_ref:
                spec_model = f"{barcode}/{internal_ref}"
            elif barcode:
                spec_model = f"{barcode}"
            elif internal_ref:  # 以防万一只有内部参考
                spec_model = f"{internal_ref}"
            else:
                spec_model = None

            # 5. 采购计量单位 -> unit
            unit = row.get('采购计量单位')

            # 6. 其他固定值或空值
            visibility_level = 0
            manual_link = ""  # 空值
            product_image = ""  # 空值 (或者根据您的JSON需求，如果是存JSON字符串可以是 '[]')
            is_enabled = True

            # 简单的非空检查 (根据数据库 NOT NULL 约束)
            if not name:
                print(f"跳过第 {index + 2} 行：名称为空")
                continue

            # --- 执行插入 SQL ---
            insert_query = """
                           INSERT INTO material_base (name, \
                                                      common_name, \
                                                      category, \
                                                      material_type, \
                                                      spec_model, \
                                                      unit, \
                                                      visibility_level, \
                                                      manual_link, \
                                                      product_image, \
                                                      is_enabled) \
                           VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s) \
                           """

            cur.execute(insert_query, (
                name,
                common_name,
                category,
                material_type,
                spec_model,
                unit,
                visibility_level,
                manual_link,
                product_image,
                is_enabled
            ))

            success_count += 1

        # 提交事务
        conn.commit()
        print(f"导入完成！成功插入 {success_count} 条数据。")

    except Exception as e:
        print(f"发生错误: {e}")
        if 'conn' in locals() and conn:
            conn.rollback()
    finally:
        if 'cur' in locals() and cur:
            cur.close()
        if 'conn' in locals() and conn:
            conn.close()


if __name__ == "__main__":
    process_excel_to_db()