large-order-fusion/utils/tools.py

# -*- coding: utf-8 -*-

import os
import re
import json
import time
from openpyxl import Workbook, load_workbook
import clickhouse_connect
from dotenv import load_dotenv
from .logger_config import setup_logger

load_dotenv()

logger = setup_logger(__name__)

# 字符串转JSON
def string_to_json(markdown_string):
    try:
        json_content = re.sub(r'^```json|\n```$', '', markdown_string, flags=re.MULTILINE).strip()
        if not json_content:
            raise ValueError("字符串中未找到有效的JSON内容")
            
        # 解析JSON内容
        json_data = json.loads(json_content)
        
        return json_data
        
    except Exception as e:
        logger.info(f"生成结果解析失败：\n{markdown_string}")

# 加载SQL文件
def load_sql_file(file_name):
    current_dir = os.path.dirname(os.path.abspath(__file__))
    sql_dir = os.path.join(current_dir, '..', 'sql')
    sql_file_path = os.path.join(sql_dir, file_name)
    
    with open(sql_file_path, 'r') as file:
        return file.read()

# 查询ClickHouse数据库
def query_clickhouse(sql_query, max_retries=3):
    """
    连接ClickHouse数据库并执行查询
    """
    for attempt in range(max_retries):
        try:
            # 创建ClickHouse客户端连接
            client = clickhouse_connect.get_client(
                host=os.getenv('HOST'),
                port=int(os.getenv('PORT')),
                user=os.getenv('USER'),
                password=os.getenv('PASSWORD'),
                database=os.getenv('DATABASE'),
                connect_timeout=int(os.getenv('CONNECT_TIMEOUT')),
                send_receive_timeout=int(os.getenv('SEND_RECEIVE_TIMEOUT'))
            )
            
            logger.info(f"第{attempt+1}次尝试连接成功")
            
            # 执行查询语句
            result = client.command(sql_query)
            return result
                
        except Exception as e:
            logger.info(f"第{attempt+1}次尝试连接失败: {str(e)}")
            if attempt < max_retries - 1:
                logger.info(f"Waiting 2 seconds before retrying...")
                time.sleep(2)
            else:
                logger.info(f"经过{max_retries}次尝试后仍然无法连接到数据库")
                return None
        finally:
            # 关闭连接
            if 'client' in locals():
                try:
                    client.disconnect()
                    logger.info("数据库连接已关闭")
                except:
                    pass

# 保存字典为JSON文件
def save_dict_to_json(data, file_path):
    try:
        with open(file_path, 'w', encoding='utf-8') as json_file:
            json.dump(data, json_file, ensure_ascii=False, indent=4)
        logger.info(f"数据已成功保存到 {file_path}")
    except Exception as e:
        logger.info(f"发生错误: {e}")

# 读取JONS文件
def read_json_file(file_path):
    try:
        with open(file_path, 'r', encoding='utf-8') as json_file:
            data = json.load(json_file)
            return data
    except FileNotFoundError:
        logger.info(f"文件 {file_path} 未找到。")
    except json.JSONDecodeError:
        logger.info(f"文件 {file_path} 不是有效的JSON格式。")
    except Exception as e:
        logger.info(f"发生错误: {e}")

# 有三种颜色
def has_three_colors(text):
    """
    判断文本中是否包含至少三种不同的颜色
    :param text: 输入文本
    :return: 如果包含至少三种不同的颜色返回 True，否则返回 False
    """
    # 定义常见颜色列表（可扩展）
    colors = ['红', '橙', '黄', '绿', '青', '蓝', '紫', '棕', '黑', '白', '灰', '金']
    
    # 使用集合来存储找到的颜色（自动去重）
    found_colors = set()
    
    # 遍历所有颜色
    for color in colors:
        # 如果颜色在文本中
        if color in text:
            # 添加到已找到的颜色集合
            found_colors.add(color)
            
            # 如果已经找到三种颜色，提前返回
            if len(found_colors) >= 3:
                return True
    
    return False

# 有四种颜色
def has_four_color(text):
    """
    判断文本中是否包含至少四种不同的颜色
    :param text: 输入文本
    :return: 如果包含至少四种不同的颜色返回 True，否则返回 False
    """
    # 定义常见颜色列表（可扩展）
    colors = ['红', '橙', '黄', '绿', '青', '蓝', '紫', '棕', '黑', '白', '灰', '金']
    
    # 使用集合来存储找到的颜色（自动去重）
    found_colors = set()
    
    # 遍历所有颜色
    for color in colors:
        # 如果颜色在文本中
        if color in text:
            # 添加到已找到的颜色集合
            found_colors.add(color)
            
            # 如果已经找到三种颜色，提前返回
            if len(found_colors) >= 4:
                return True
    
    return False

def read_excel(file_path, sheet_name=None, sku_column='A'):
    # 加载工作簿
    wb = load_workbook(filename=file_path)
    
    # 获取工作表
    if sheet_name:
        ws = wb[sheet_name]
    else:
        ws = wb.active
    
    # 读取 SKU 数据
    skus = []
    for row in ws.iter_rows(min_row=2, values_only=True):  # 假设第一行是表头
        sku = row[ord(sku_column.upper()) - ord('A')]  # 将列字母转换为索引
        if sku:  # 只添加非空值
            skus.append(str(sku))
    
    return skus

def extract_fields(dict_list, fields, default=None):
    """
    提取字典列表中每个字典的指定字段
    
    参数:
        dict_list (list): 包含字典的列表
        fields (list): 需要提取的字段名称列表
        default: 当字段不存在时的默认值，默认为None
    
    返回:
        list: 只包含指定字段的新字典列表
    """
    result = []
    for d in dict_list:
        # 对每个字典，提取指定字段，不存在则用default填充
        extracted = {field: d.get(field, default) for field in fields}
        result.append(extracted)
    return result

def create_sku_excel(filename="aaa.xlsx", sheet_name="SKU列表"):
    """
    创建一个只有SKU字段的Excel文件
    :param filename: 文件名（默认aaa.xlsx）
    :param sheet_name: 工作表名（默认SKU列表）
    :return: None
    """
    wb = Workbook()
    ws = wb.active
    ws.title = sheet_name  # 设置工作表名称
    ws['A1'] = "SKU"       # 写入字段名
    wb.save(filename)
    logger.info(f"文件 '{filename}' 创建成功，包含SKC字段")

def add_sku_to_excel(sku_value, filename="aaa.xlsx"):
    """
    向Excel文件追加SKU数据
    :param sku_value: 要添加的SKU值（如11DAAF35425）
    :param filename: 文件名（默认aaa.xlsx）
    :return: None
    """
    # 检查文件是否存在
    if not os.path.exists(filename):
        logger.info(f"错误：文件 '{filename}' 未找到！")
        return

    try:
        wb = load_workbook(filename)
        ws = wb.active
        
        # 找到第一个空行（兼容有/无表头的情况）
        next_row = ws.max_row + 1 if ws['A1'].value is not None else 1
        
        # 如果是空文件，自动添加SKU表头
        if next_row == 1:
            ws['A1'] = "SKC"
            next_row = 2
        
        ws[f'A{next_row}'] = sku_value
        wb.save(filename)
        logger.info(f"SKU '{sku_value}' 已成功添加到文件 '{filename}' 的第 {next_row} 行")
    except Exception as e:
        logger.info(f"操作失败：{str(e)}")

def add_suffix(filepath, suffix="_failed"):
    """
    在文件名后添加'_failed'后缀，保留原扩展名
    :param filepath: 原始文件绝对路径（如C:\data\example.xlsx）
    :return: 添加后缀后的新路径（如C:\data\example_failed.xlsx）
    """
    # 分离目录、文件名和扩展名
    dirname = os.path.dirname(filepath)
    basename, ext = os.path.splitext(os.path.basename(filepath))
    
    # 构建新文件名
    new_basename = f"{basename}{suffix}{ext}"
    new_path = os.path.join(dirname, new_basename)
    
    return new_path

if __name__ == '__main__':

    # sql_query = load_sql_file('get_sku_info.sql')

    # sql_query = sql_query.replace('goods_id_holder', '2739585094554112')
    # sql_query = sql_query.replace('color_id_holder', '2055025368795719')
    # sql_result = query_clickhouse(sql_query)

    # logger.info(sql_result)


    # sql_query = load_sql_file('get_sku_id.sql')

    # sql_query = sql_query.replace('g_code_holder', '1E5C6M200')
    # sql_query = sql_query.replace('cl_code_holder', '91Y')
    # sql_result = query_clickhouse(sql_query)

    # logger.info(sql_result)

    logger.info(read_excel("./data/primary_sku/主推款1015.xlsx"))