DB-GPT/dbgpt/app/scene/chat_data/chat_excel/excel_reader.py

import io
import json
import logging
import os

import chardet
import duckdb
import numpy as np
import pandas as pd
import sqlparse
from pyparsing import (
    CaselessKeyword,
    Forward,
    Literal,
    Optional,
    Regex,
    Word,
    alphanums,
    delimitedList,
)

from dbgpt.util.file_client import FileClient
from dbgpt.util.pd_utils import csv_colunm_foramt
from dbgpt.util.string_utils import is_chinese_include_number

logger = logging.getLogger(__name__)


def excel_colunm_format(old_name: str) -> str:
    new_column = old_name.strip()
    new_column = new_column.replace(" ", "_")
    return new_column


def detect_encoding(file_path):
    # 读取文件的二进制数据
    with open(file_path, "rb") as f:
        data = f.read()
    # 使用 chardet 来检测文件编码
    result = chardet.detect(data)
    encoding = result["encoding"]
    confidence = result["confidence"]
    return encoding, confidence


def add_quotes_ex(sql: str, column_names):
    sql = sql.replace("`", '"')
    for column_name in column_names:
        if sql.find(column_name) != -1 and sql.find(f'"{column_name}"') == -1:
            sql = sql.replace(column_name, f'"{column_name}"')
    return sql


def parse_sql(sql):
    # 定义关键字和标识符
    select_stmt = Forward()
    column = Regex(r"[\w一-龥]*")
    table = Word(alphanums)
    join_expr = Forward()
    where_expr = Forward()
    group_by_expr = Forward()
    order_by_expr = Forward()

    select_keyword = CaselessKeyword("SELECT")
    from_keyword = CaselessKeyword("FROM")
    join_keyword = CaselessKeyword("JOIN")
    on_keyword = CaselessKeyword("ON")
    where_keyword = CaselessKeyword("WHERE")
    group_by_keyword = CaselessKeyword("GROUP BY")
    order_by_keyword = CaselessKeyword("ORDER BY")
    and_keyword = CaselessKeyword("AND")
    or_keyword = CaselessKeyword("OR")
    in_keyword = CaselessKeyword("IN")
    not_in_keyword = CaselessKeyword("NOT IN")

    # 定义语法规则
    select_stmt <<= (
        select_keyword
        + delimitedList(column)
        + from_keyword
        + delimitedList(table)
        + Optional(join_expr)
        + Optional(where_keyword + where_expr)
        + Optional(group_by_keyword + group_by_expr)
        + Optional(order_by_keyword + order_by_expr)
    )

    join_expr <<= join_keyword + table + on_keyword + column + Literal("=") + column

    where_expr <<= (
        column + Literal("=") + Word(alphanums) + Optional(and_keyword + where_expr)
        | column + Literal(">") + Word(alphanums) + Optional(and_keyword + where_expr)
        | column + Literal("<") + Word(alphanums) + Optional(and_keyword + where_expr)
    )

    group_by_expr <<= delimitedList(column)

    order_by_expr <<= column + Optional(Literal("ASC") | Literal("DESC"))

    # 解析 SQL 语句
    parsed_result = select_stmt.parseString(sql)

    return parsed_result.asList()


def add_quotes(sql, column_names=[]):
    sql = sql.replace("`", "")
    sql = sql.replace("'", "")
    parsed = sqlparse.parse(sql)
    for stmt in parsed:
        for token in stmt.tokens:
            deep_quotes(token, column_names)
    return str(parsed[0])


def deep_quotes(token, column_names=[]):
    if hasattr(token, "tokens"):
        for token_child in token.tokens:
            deep_quotes(token_child, column_names)
    else:
        if is_chinese_include_number(token.value):
            new_value = token.value.replace("`", "").replace("'", "")
            token.value = f'"{new_value}"'


def get_select_clause(sql):
    parsed = sqlparse.parse(sql)[0]  # 解析 SQL 语句，获取第一个语句块

    select_tokens = []
    is_select = False

    for token in parsed.tokens:
        if token.is_keyword and token.value.upper() == "SELECT":
            is_select = True
        elif is_select:
            if token.is_keyword and token.value.upper() == "FROM":
                break
            select_tokens.append(token)
    return "".join(str(token) for token in select_tokens)


def parse_select_fields(sql):
    parsed = sqlparse.parse(sql)[0]  # 解析 SQL 语句，获取第一个语句块
    fields = []

    for token in parsed.tokens:
        # 使用 flatten() 方法合并 '2022' 和 '年' 为一个 token
        if token.match(sqlparse.tokens.Literal.String.Single):
            token.flatten()
        if isinstance(token, sqlparse.sql.Identifier):
            fields.append(token.get_real_name())

    # 处理中文
    fields = [field.replace(f"field", f'"{field}"') for field in fields]

    return fields


def add_quotes_to_chinese_columns(sql, column_names=[]):
    parsed = sqlparse.parse(sql)
    for stmt in parsed:
        process_statement(stmt, column_names)
    return str(parsed[0])


def process_statement(statement, column_names=[]):
    if isinstance(statement, sqlparse.sql.IdentifierList):
        for identifier in statement.get_identifiers():
            process_identifier(identifier)
    elif isinstance(statement, sqlparse.sql.Identifier):
        process_identifier(statement, column_names)
    elif isinstance(statement, sqlparse.sql.TokenList):
        for item in statement.tokens:
            process_statement(item)


def process_identifier(identifier, column_names=[]):
    # if identifier.has_alias():
    #     alias = identifier.get_alias()
    #     identifier.tokens[-1].value = '[' + alias + ']'
    if hasattr(identifier, "tokens") and identifier.value in column_names:
        if is_chinese(identifier.value):
            new_value = get_new_value(identifier.value)
            identifier.value = new_value
            identifier.normalized = new_value
            identifier.tokens = [sqlparse.sql.Token(sqlparse.tokens.Name, new_value)]
    else:
        if hasattr(identifier, "tokens"):
            for token in identifier.tokens:
                if isinstance(token, sqlparse.sql.Function):
                    process_function(token)
                elif token.ttype in sqlparse.tokens.Name:
                    new_value = get_new_value(token.value)
                    token.value = new_value
                    token.normalized = new_value
                elif token.value in column_names:
                    new_value = get_new_value(token.value)
                    token.value = new_value
                    token.normalized = new_value
                    token.tokens = [sqlparse.sql.Token(sqlparse.tokens.Name, new_value)]


def get_new_value(value):
    return f""" "{value.replace("`", "").replace("'", "").replace('"', "")}" """


def process_function(function):
    function_params = list(function.get_parameters())
    # for param in function_params:
    for i in range(len(function_params)):
        param = function_params[i]
        # 如果参数部分是一个标识符（字段名）
        if isinstance(param, sqlparse.sql.Identifier):
            # 判断是否需要替换字段值
            # if is_chinese(param.value):
            # 替换字段值
            new_value = get_new_value(param.value)
            # new_parameter = sqlparse.sql.Identifier(f'[{param.value}]')
            function_params[i].tokens = [
                sqlparse.sql.Token(sqlparse.tokens.Name, new_value)
            ]
    print(str(function))


def is_chinese(text):
    for char in text:
        if "一" <= char <= "鿿":
            return True
    return False


class ExcelReader:
    def __init__(self, conv_uid: str, file_param: str):
        self.conv_uid = conv_uid
        self.file_param = file_param
        if isinstance(file_param, str) and os.path.isabs(file_param):
            file_name = os.path.basename(file_param)
            self.file_name_without_extension = os.path.splitext(file_name)[0]
            encoding, confidence = detect_encoding(file_param)

            self.excel_file_name = file_name
            self.extension = os.path.splitext(file_name)[1]

            file_info = file_param
        else:
            if isinstance(file_param, dict):
                file_path = file_param.get("file_path", None)
                if not file_path:
                    raise ValueError("Not find file path!")
                else:
                    file_name = os.path.basename(file_path.replace(f"{conv_uid}_", ""))

            else:
                temp_obj = json.loads(file_param)
                file_path = temp_obj.get("file_path", None)
                file_name = os.path.basename(file_path.replace(f"{conv_uid}_", ""))

            self.file_name_without_extension = os.path.splitext(file_name)[0]

            self.excel_file_name = file_name
            self.extension = os.path.splitext(file_name)[1]

            file_client = FileClient()
            file_info = file_client.read_file(
                conv_uid=self.conv_uid, file_key=file_path
            )

            result = chardet.detect(file_info)
            encoding = result["encoding"]
            confidence = result["confidence"]

        logger.info(
            f"File Info:{len(file_info)},Detected Encoding: {encoding} (Confidence: {confidence})"
        )

        # read excel file
        if file_name.endswith(".xlsx") or file_name.endswith(".xls"):
            df_tmp = pd.read_excel(file_info, index_col=False)
            self.df = pd.read_excel(
                file_info,
                index_col=False,
                converters={i: csv_colunm_foramt for i in range(df_tmp.shape[1])},
            )
        elif file_name.endswith(".csv"):
            df_tmp = pd.read_csv(
                file_info if isinstance(file_info, str) else io.BytesIO(file_info),
                index_col=False,
                encoding=encoding,
            )
            self.df = pd.read_csv(
                file_info if isinstance(file_info, str) else io.BytesIO(file_info),
                index_col=False,
                encoding=encoding,
                converters={i: csv_colunm_foramt for i in range(df_tmp.shape[1])},
            )
        else:
            raise ValueError("Unsupported file format.")

        self.df.replace("", np.nan, inplace=True)

        # 修改的部分

        unnamed_columns_tmp = [
            col
            for col in df_tmp.columns
            if col.startswith("Unnamed") and df_tmp[col].isnull().all()
        ]
        df_tmp.drop(columns=unnamed_columns_tmp, inplace=True)

        self.df = self.df[df_tmp.columns.values]
        #

        self.columns_map = {}
        for column_name in df_tmp.columns:
            self.df[column_name] = self.df[column_name].astype(str)
            self.columns_map.update({column_name: excel_colunm_format(column_name)})
            try:
                self.df[column_name] = pd.to_datetime(self.df[column_name]).dt.strftime(
                    "%Y-%m-%d"
                )
            except ValueError:
                try:
                    self.df[column_name] = pd.to_numeric(self.df[column_name])
                except ValueError:
                    try:
                        self.df[column_name] = self.df[column_name].astype(str)
                    except Exception:
                        print("Can't transform column: " + column_name)

        self.df = self.df.rename(columns=lambda x: x.strip().replace(" ", "_"))

        # connect DuckDB
        self.db = duckdb.connect(database=":memory:", read_only=False)

        self.table_name = "excel_data"
        # write data in duckdb
        self.db.register(self.table_name, self.df)

        # 获取结果并打印表结构信息
        result = self.db.execute(f"DESCRIBE {self.table_name}")
        columns = result.fetchall()
        for column in columns:
            print(column)

    def run(self, sql):
        try:
            if f'"{self.table_name}"' in sql:
                sql = sql.replace(f'"{self.table_name}"', self.table_name)
            sql = add_quotes_to_chinese_columns(sql)
            print(f"excute sql:{sql}")
            results = self.db.execute(sql)
            colunms = []
            for descrip in results.description:
                colunms.append(descrip[0])
            return colunms, results.fetchall()
        except Exception as e:
            logger.error(f"excel sql run error!, {str(e)}")
            raise ValueError(f"Data Query Exception!\\nSQL[{sql}].\\nError:{str(e)}")

    def get_df_by_sql_ex(self, sql):
        colunms, values = self.run(sql)
        return pd.DataFrame(values, columns=colunms)

    def get_sample_data(self):
        return self.run(f"SELECT * FROM {self.table_name} LIMIT 5;")