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}"
    "?ssl_verify_cert=false&ssl_verify_identity=false"
)
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
from dotenv import load_dotenv
load_dotenv()
MONGO_URI = os.getenv("MONGO_URI", "")
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.5)

#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 = {
        "payments": prompt_detail_table.prompt_payments,
        "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
    print("SQL query:", sql_query)
    parsed_query = sqlglot.parse_one(sql_query, read="mysql")
    
    # 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,user_id: int) -> 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à:
User Id: {user_id}
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).
Loại bỏ các câu comment chỉ trả lại câu SQL chính xác sau khi bạn đã chỉnh sửa.
Vui lòng chỉnh sửa một cách chính xác nhất.
Trả lời dưới dạng một truy vấn SQL duy nhất.
** Tham khảo thêm các lỗi sau để tránh:
 - "- Tránh các lỗi như :\n"
    "  (1054, \"Unknown column 'oi.pro_id' in 'field list'\")\n . Luôn đảm bảo bạn không bao giờ bị lỗi này"
    "  (1054, \"Unknown column 'oi.note' in 'field list'\") . Luôn đảm bảo bạn không bao giờ bị lỗi này\n"
    "  (1054, \"Unknown column 'oi.size' in 'field list'\") . Luôn đảm bảo bạn không bao giờ bị lỗi này \n"
    "  (1054, \"Unknown column 'c.is_deleted' in 'on clause'\").  Luôn đảm bảo bạn không bao giờ bị lỗi này\n"
"""
 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 regenerate_sql_until_safe():
     max_retry = 5
     retry_count = 0

     while retry_count < max_retry:
        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:
                regenerated_sql = clean_sql(regenerated_sql)
                if not re.search(r"%{1,2}s", regenerated_sql):  # đã sạch
                    return regenerated_sql
            retry_count += 1
        except Exception as e:
            return f"❌ Lỗi khi tạo lại truy vấn lần {retry_count + 1}: {str(e)}"

     return "❌ Lỗi: Không thể tạo được truy vấn an toàn sau nhiều lần thử."


    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:
        print("Clean sql1:", sql)
        if isinstance(sql, dict) and sql:
           first_value = next(iter(sql.values()))
           sql = first_value
        sql = re.sub(r"```sql", "", sql, flags=re.IGNORECASE)
        sql = sql.replace("```sql", "")
        sql = re.sub(r"```", "", sql)
        return sql.strip()
        # 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):
     match = re.search(r"\[SQL:\s*```sql\s*(.*?)\s*```]", data, re.DOTALL)
     if match:
        return clean_sql(match.group(1))
     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):
     print("NHận sql: ",sql )
     data = sql
     if isinstance(data, dict) and data:
      print("Là dict")
      first_key, first_value = next(iter(data.items()))
      sql = first_value
     elif isinstance(data, list) and data:
       sql = "\n\n".join(data)

     sql_clean = clean_sql(sql)
     print("SQL Clean: ", sql_clean)
     if re.search(r"%{1,2}s", sql_clean):
       regenerated_sql = regenerate_sql_until_safe()
       sql_clean = clean_sql(regenerated_sql)
     result = get_schemas_from_sql(sql_clean, schema_mapping)
     prompt = build_sql_fix_prompt(schemas_result=result,sql = str(sql_clean),user_id =user_id)
     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)


    
    PROMPT_CUSTOM.template += (
        "\n\n⚠️ Lưu ý: Do NOT use IF...THEN...ELSE...END in SQL. "
        "Only use plain SELECT, INSERT, UPDATE, JOIN, SET statements. "
        "Tuyệt đối cấm dùng %%s hoặc %s để lấy giá trị từ biến. "
        "SQL phải chứa giá trị cố định, không dùng placeholder hay biến bind."
    )

    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)}"