Spaces:
Sleeping
Sleeping
File size: 2,264 Bytes
6bd3e57 e252f82 6bd3e57 e252f82 6bd3e57 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | """
Database utilities for querying and writing to MySQL through an SSH tunnel.
"""
from dotenv import load_dotenv
load_dotenv()
import pandas as pd
from sqlalchemy import text
from .db_connections import connect_ssh_tunnel, create_db_engine, disconnect
def read_sql(query_str: str) -> pd.DataFrame:
"""
Execute a read SQL query and return the results as a pandas DataFrame.
Uses pandas built-in read_sql_query method.
Args:
query_str: SQL query string to execute
Returns:
pandas DataFrame with query results
Raises:
Exception: Re-raises any database or connection errors for proper error handling
"""
tunnel = None
engine = None
try:
tunnel = connect_ssh_tunnel()
engine = create_db_engine(tunnel)
df = pd.read_sql_query(sql=text(query_str), con=engine)
return df
except Exception as e:
print(f"Error executing read query: {e}")
# Re-raise the exception so it can be handled by the calling function
raise e
finally:
disconnect(engine, tunnel)
def write_to_table(df: pd.DataFrame, table_name: str, if_exists: str = 'append', schema: str = None) -> bool:
"""
Write a pandas DataFrame to a MySQL table using pandas built-in to_sql method.
Args:
df: pandas DataFrame to write
table_name: Name of the target table
if_exists: What to do if table exists ('append', 'replace', 'fail')
schema: Database schema name (optional)
Returns:
True if successful, False otherwise
"""
tunnel = None
engine = None
try:
tunnel = connect_ssh_tunnel()
engine = create_db_engine(tunnel)
df.to_sql(
name=table_name,
con=engine,
if_exists=if_exists,
index=False,
schema=schema,
method='multi',
chunksize=1000
)
print(f"Successfully wrote {len(df)} rows to table '{table_name}' (mode: {if_exists})")
return True
except Exception as e:
print(f"Error writing DataFrame to table: {e}")
return False
finally:
disconnect(engine, tunnel)
|