function/analyze/execute_query.py

import sys
import os

# Thêm thư mục gốc (chứa thư mục "function") vào sys.path
sys.path.append(os.path.abspath(os.path.join(os.getcwd(), "..")))
import os
from dotenv import load_dotenv

# Load biến môi trường từ file .env
load_dotenv()

DB_HOST = os.getenv("DB_HOST")
DB_USER = os.getenv("DB_USER")
DB_PASSWORD = os.getenv("DB_PASSWORD")
DB_NAME =  os.getenv("DB_NAME")
DB_PORT =  os.getenv("DB_PORT")

import os
from urllib.parse import quote

password = os.getenv("DB_PASSWORD")  # VD: 'Yahana0509@'
DB_PASSWORD = quote(password)
# Tạo connection string
connection_uri = f"mysql+pymysql://{DB_USER}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_NAME}"
from langchain_community.utilities.sql_database import SQLDatabase
from langchain_experimental.sql import SQLDatabaseChain
import sys
import os
import pymysql
from fastapi import HTTPException
import sys
import os


from fastapi.encoders import jsonable_encoder
import re
from function.prompt import prompt_main as prompt
import function.prompt.prompt_custom as prompt_cus
os.environ["GOOGLE_API_KEY"] = "AIzaSyCO-RlqYewC4e9BEPoC8m-AxHUY7J3_o2E"
from bson import ObjectId
db = SQLDatabase.from_uri(connection_uri)
from dotenv import load_dotenv
import filter.filter_role as filter_role_1
import filter.filter_sql_injection as filter_sql_injection_1
import filter.result as query_result_1
import response.ResponseChat as res_chat
from datetime import datetime
import pytz
from mongoengine import connect
import os
import nltk
import function.agent.pipeline_agent as pipeline_agent
nltk.download('punkt')
from models.Database_Entity import User, ChatHistory, DetailChat
MONGO_URI = "mongodb+srv://kltn20133118:XEYSVzYvEwsp5Cvo@cluster0.nnsw9.mongodb.net/?retryWrites=true&w=majority&appName=Cluster0"
connect("chatbot_hmdrinks", host=MONGO_URI)
import re

def contains_delete(sql: str) -> bool:
    return bool(re.search(r'\bdelete\b', sql, re.IGNORECASE))
load_dotenv()

#setup model
from bson import ObjectId
import random
from langchain_google_genai import GoogleGenerativeAIEmbeddings, ChatGoogleGenerativeAI
llm1 = ChatGoogleGenerativeAI(model='gemini-2.5-flash-preview-05-20',temperature=0.6)

#setup kết nối database
db = SQLDatabase.from_uri(connection_uri)
db_chain = SQLDatabaseChain.from_llm(llm=llm1,db=db,prompt= prompt.PROMPT)

from prompt.prompt_syntax_insert import is_insert_related_to_product_category_variant, filter_syntax_sql
import sqlparse


import sys
import os
sys.path.append(os.path.abspath(os.path.join(os.getcwd(), '..')))
from prompt import prompt_detail_table

schema_mapping = {
        "user": prompt_detail_table.prompt_users,
        "user_voucher": prompt_detail_table.prompt_user_voucher,
        "category": prompt_detail_table.prompt_categort,
        "category_translation": prompt_detail_table.prompt_category_translation,
        "cart": prompt_detail_table.prompt_cart,
        "cart_item": prompt_detail_table.prompt_cart_item,
        "orders":prompt_detail_table.prompt_orders,
        "order_item": prompt_detail_table.prompt_order_item,
        "payment": prompt_detail_table.prompt_payments,
        "favourite": prompt_detail_table.prompt_favourite,
        "favourite_item": prompt_detail_table.prompt_fav_item,
        "post": prompt_detail_table.prompt_post,
        "post_translation": prompt_detail_table.prompt_post_translation,
        "product": prompt_detail_table.prompt_product,
        "product_translation": prompt_detail_table.prompt_product_translation,
        "shipment": prompt_detail_table.prompt_shipment,
        "product_variants": prompt_detail_table.prompt_product_variants,
        "review": prompt_detail_table.prompt_review,
        "user_coin": prompt_detail_table.prompt_user_coin,
        "absence": prompt_detail_table.prompt_absence,
        "cart_group": prompt_detail_table.prompt_cart_group,
        "cart_item_group": prompt_detail_table.prompt_cartitem_group,
        "group_orders": prompt_detail_table.prompt_group_orders,
        "payments_group": prompt_detail_table.prompt_payments_group,
        "group_order_members":prompt_detail_table.prompt_group_orders_member,
        "shipment_group":prompt_detail_table.prompt_shipment_group,
        "shipper_attendance":prompt_detail_table.prompt_shipper_attendance,
        "shipper_commission_detail":prompt_detail_table.prompt_shipper_commission_detail,
        "shipper_salary_summary":prompt_detail_table.prompt_shipper_salary_summary,
        "voucher": prompt_detail_table.prompt_voucher
    }


def get_schemas_from_sql(sql_query: str, schema_mapping: dict):
    import sqlglot

    parsed_query = sqlglot.parse_one(sql_query, read="sqlite")
    
    # Lấy danh sách bảng duy nhất trong query
    table_names = list({t.name for t in parsed_query.find_all(sqlglot.exp.Table)})

    schemas_used = {}
    for table in table_names:
        if table in schema_mapping:
            schemas_used[table] = schema_mapping[table]
        else:
            print(f"⚠️ Warning: Table '{table}' not found in schema_mapping")

    # Gom toàn bộ schema thành 1 chuỗi duy nhất
    all_schemas = "\n\n".join(
        [f"Schema for table '{table}':\n{schemas_used[table]}" for table in schemas_used]
    )

    return all_schemas


def build_sql_fix_prompt(schemas_result: dict, sql: str) -> str:
   
 prompt = f"""
Bạn là một chuyên gia cơ sở dữ liệu.

Dưới đây là mô tả schema chi tiết của các bảng có trong hệ thống:

{schemas_result}

---

Dưới đây là một câu SQL đang bị lỗi do không đúng tên bảng hoặc tên cột:

```sql
{sql.strip()}
Yêu cầu của bạn là:

Dựa trên các schema ở trên, hãy kiểm tra và chỉnh sửa câu SQL sao cho:
Tên bảng, tên cột phải chính xác theo schema.
Logic và mục đích của truy vấn được giữ nguyên.
Chỉ trả lại phần SQL đã được chỉnh sửa (không giải thích, không chú thích, không thêm nhận xét).
Trả lời dưới dạng một truy vấn SQL duy nhất.
""".strip()
 return prompt

async def execute_query_user(user_input: str, user_id: int, languages: str, role: str):
    PROMPT_CUSTOM = await prompt_cus.get_prompt_custom(user_input)
    check_insert = is_insert_related_to_product_category_variant(user_input)

    db_config = {
        "host": os.getenv("DB_HOST"),
        "user": os.getenv("DB_USER"),
        "database": os.getenv("DB_NAME"),
        "password": os.getenv("DB_PASSWORD"),
        "port": int(os.getenv("DB_PORT", 3306)),
        "charset": "utf8mb4",
        "cursorclass": pymysql.cursors.DictCursor,
    }

    def execute_query_with_pymysql(query, multi=False):
        connection = pymysql.connect(**db_config)
        try:
            with connection.cursor() as cursor:
                results = []
                if multi:
                    # Dùng sqlparse để chia câu truy vấn cho an toàn
                    statements = sqlparse.split(query)
                    for stmt in statements:
                        stmt = stmt.strip()
                        if stmt:
                            cursor.execute(stmt)
                            results.append(cursor.fetchall())
                else:
                    cursor.execute(query)
                    results = cursor.fetchall()
                connection.commit()
                return results
        except pymysql.MySQLError as e:
            return str(e)
        finally:
            connection.close()

    def clean_sql(sql: str) -> str:
        # Xoá các markdown syntax kiểu ```sql ... ```
        sql = re.sub(r"```sql", "", sql, flags=re.IGNORECASE)
        sql = sql.replace("```sql", "")
        # sql = re.sub(r'%%s', r'%s', sql)
        sql = re.sub(r"```", "", sql)
        return sql.strip()

    def extract_sql_from_response(data):
    # Định dạng markdown block [SQL: ```sql ... ```]
     match = re.search(r"\[SQL:\s*```sql\s*(.*?)\s*```]", data, re.DOTALL)
     if match:
        return clean_sql(match.group(1))
    
    # Định dạng đơn giản SQLQuery: ...
     match = re.search(r"SQLQuery:\s*(.*)", data, re.DOTALL)
     if match:
        return clean_sql(match.group(1))

     return None

    def extract_sql_from_error(error_msg):
    # Case 1: [SQL: ```sql ... ```]
     match = re.search(r"\[SQL:\s*```sql\s*(.*?)\s*```]", error_msg, re.DOTALL)
     if match:
        return clean_sql(match.group(1))


     match = re.search(r"```(?:sql)?\s*\r?\n(.*?)```", error_msg, re.DOTALL)
     if match:
        return clean_sql(match.group(1))

     return None

    def process_and_execute_sql(sql):
     sql_clean = clean_sql(sql)
     result = get_schemas_from_sql(sql_clean, schema_mapping)
     prompt = build_sql_fix_prompt(result,sql =sql_clean)
     from advance_shopping.call_gemini import tool_call
     data = tool_call.generate(prompt = prompt)
     sql_clean = clean_sql(data)
     print("SQL step2: ", sql_clean)
     if contains_delete(sql_clean):
            return "Lỗi: Bạn không dược phép thực hiện truy vấn DELETE trong hệ thống này."

     if re.search(r"\bIF\b.*\bTHEN\b", sql_clean, re.IGNORECASE):
        return "❌ Lỗi: Không được dùng IF...THEN trong SQL. Vui lòng chia nhỏ truy vấn."

     if check_insert:
        check_syntax = filter_syntax_sql(sql_clean, PROMPT_CUSTOM, user_input)
        if check_syntax is True:
            # Tách từng câu và thực thi tuần tự
            try:
                connection = pymysql.connect(**db_config)
                with connection.cursor() as cursor:
                    statements = sqlparse.split(sql_clean)
                    results = []
                    for stmt in statements:
                        stmt = stmt.strip()
                        if stmt:
                            cursor.execute(stmt)
                            try:
                                results.append(cursor.fetchall())
                            except:
                                results.append("✅ OK")  # Có thể là câu SET hoặc INSERT
                    connection.commit()
                    return results
            except Exception as e:
                return f"❌ Lỗi khi thực thi từng truy vấn: {str(e)}"
            finally:
                connection.close()
        else:
            try:
                regenerated_data = db_chain.run(f"""
                    Role: {text_role}
                    Language: {languages}
                    Question: {user_input}.
                """)
                regenerated_sql = extract_sql_from_response(regenerated_data)
                if regenerated_sql:
                    return process_and_execute_sql(regenerated_sql)
                else:
                    return "❌ Lỗi: Không thể tạo lại truy vấn hợp lệ."
            except Exception as regen_error:
                return f"❌ Lỗi khi tạo lại truy vấn: {str(regen_error)}"
     else:
        return execute_query_with_pymysql(sql_clean, multi=True)


    if "Do not use IF...THEN" not in PROMPT_CUSTOM.template:
        PROMPT_CUSTOM.template += (
            "\n\n⚠️ Luu ý: Do NOT use IF...THEN...ELSE...END in SQL. "
            "Only use plain SELECT, INSERT, UPDATE, JOIN, SET statements. Không dùng %%s để lấy giá trị từ biến. "
        )

    db_chain = SQLDatabaseChain.from_llm(llm=llm1, db=db, prompt=PROMPT_CUSTOM)
    text_role = f"{role} (userId = {user_id})" if role == "ADMIN" else f"{role} (userId = {user_id}), not role ADMIN"

    try:
        data = db_chain.run(f"""
            Role: {text_role}
            Language: {languages}
            Question: {user_input}.
        """)
        extracted_sql = extract_sql_from_response(data)
        if extracted_sql:
            return process_and_execute_sql(extracted_sql)
        else:
            return data

    except Exception as e:
         error_message = str(e)
         extracted_sql = extract_sql_from_error(error_message)

         if extracted_sql:
            fix_sql = re.sub(r"```sql", "", extracted_sql)
            fix_sql = extracted_sql.replace("```", "")
            fix_sql = extracted_sql.replace("```sql", "")
            # fix_sql = re.sub(r'%%s', r'%s', fix_sql)
            if contains_delete(fix_sql):
               return "Lỗi: Bạn không dược phép thực hiện truy vấn DELETE trong hệ thống này."
            return process_and_execute_sql(fix_sql)
         else:
            # Không trích xuất được SQL -> Thử tạo lại truy vấn từ đầu
            try:
                regenerated_data = db_chain.run(f"""
                    Role: {text_role}
                    Language: {languages}
                    Question: {user_input}.
                """)
                regenerated_sql = extract_sql_from_response(regenerated_data)
                if regenerated_sql:
                    return process_and_execute_sql(regenerated_sql)
                else:
                    return "❌ Lỗi: Không thể tạo lại truy vấn hợp lệ sau lỗi."
            except Exception as regen_error:
                return f"❌ Lỗi khi tạo lại truy vấn từ lỗi: {str(regen_error)}"