100 lines
2.8 KiB
Python
100 lines
2.8 KiB
Python
|
import os
|
||
|
|
import json
|
||
|
|
import pymysql
|
||
|
|
import pandas as pd
|
||
|
|
from typing import Dict, List, Optional, Any
|
||
|
|
from dotenv import load_dotenv
|
||
|
|
import re
|
||
|
|
|
||
|
|
# Load environment variables
|
||
|
|
load_dotenv(override=True)
|
||
|
|
|
||
|
|
# Database configuration
|
||
|
|
DB_CONFIG = {
|
||
|
|
"host": os.getenv("DB_HOST", "localhost"),
|
||
|
|
"port": int(os.getenv("DB_PORT", "9030")),
|
||
|
|
"user": os.getenv("DB_USER", "root"),
|
||
|
|
"password": os.getenv("DB_PASSWORD", ""),
|
||
|
|
"database": os.getenv("DB_DATABASE", ""),
|
||
|
|
"charset": "utf8mb4",
|
||
|
|
"cursorclass": pymysql.cursors.DictCursor
|
||
|
|
}
|
||
|
|
|
||
|
|
def get_db_connection(db_name: Optional[str] = None):
|
||
|
|
"""
|
||
|
|
Get database connection
|
||
|
|
|
||
|
|
Args:
|
||
|
|
db_name: Specify the database name to connect to, use default config if None
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Database connection
|
||
|
|
"""
|
||
|
|
if db_name:
|
||
|
|
# Use default config but override database name
|
||
|
|
config = DB_CONFIG.copy()
|
||
|
|
config["database"] = db_name
|
||
|
|
return pymysql.connect(**config)
|
||
|
|
else:
|
||
|
|
# Use default config
|
||
|
|
return pymysql.connect(**DB_CONFIG)
|
||
|
|
|
||
|
|
def get_db_name() -> str:
|
||
|
|
"""Get the currently configured default database name"""
|
||
|
|
return DB_CONFIG["database"] or os.getenv("DB_DATABASE", "")
|
||
|
|
|
||
|
|
def execute_query(sql, db_name: Optional[str] = None):
|
||
|
|
"""
|
||
|
|
Execute SQL query and return results
|
||
|
|
|
||
|
|
Args:
|
||
|
|
sql: SQL query statement
|
||
|
|
db_name: Specify the database name to connect to, use default config if None
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Query results
|
||
|
|
"""
|
||
|
|
conn = get_db_connection(db_name)
|
||
|
|
try:
|
||
|
|
with conn.cursor() as cursor:
|
||
|
|
# Set connection character set to utf8 before executing query
|
||
|
|
cursor.execute("SET NAMES utf8")
|
||
|
|
|
||
|
|
# Execute the actual query
|
||
|
|
cursor.execute(sql)
|
||
|
|
result = cursor.fetchall()
|
||
|
|
return result
|
||
|
|
finally:
|
||
|
|
conn.close()
|
||
|
|
|
||
|
|
def execute_query_df(sql, db_name: Optional[str] = None):
|
||
|
|
"""
|
||
|
|
Execute SQL query and return pandas DataFrame
|
||
|
|
|
||
|
|
Args:
|
||
|
|
sql: SQL query statement
|
||
|
|
db_name: Specify the database name to connect to, use default config if None
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
pandas DataFrame
|
||
|
|
"""
|
||
|
|
conn = get_db_connection(db_name)
|
||
|
|
try:
|
||
|
|
# Use a temporary cursor to execute the query and get results
|
||
|
|
with conn.cursor() as cursor:
|
||
|
|
# Set connection character set to utf8 before executing query
|
||
|
|
cursor.execute("SET NAMES utf8")
|
||
|
|
|
||
|
|
# Execute the actual query
|
||
|
|
cursor.execute(sql)
|
||
|
|
result = cursor.fetchall()
|
||
|
|
|
||
|
|
# If no results, return empty DataFrame
|
||
|
|
if not result:
|
||
|
|
return pd.DataFrame()
|
||
|
|
|
||
|
|
# Manually convert dict results to DataFrame
|
||
|
|
df = pd.DataFrame(result)
|
||
|
|
return df
|
||
|
|
finally:
|
||
|
|
conn.close()
|