Skip to content

Commit 83684c5

Browse files
kks32kks32
committed
DesignSafe database
1 parent 44ef9b6 commit 83684c5

File tree

6 files changed

+382
-44
lines changed

6 files changed

+382
-44
lines changed

‎dapi/client.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
from . import apps as apps_module
55
from . import files as files_module
66
from . import jobs as jobs_module
7+
from .db.accessor import DatabaseAccessor
78

89
# Import only the necessary classes/functions from jobs
910
from .jobs import SubmittedJob # JobDefinition is no longer needed
@@ -26,9 +27,12 @@ def __init__(self, tapis_client: Optional[Tapis] = None, **auth_kwargs):
2627
self.tapis = tapis_client
2728
else:
2829
self.tapis = auth.init(**auth_kwargs)
30+
31+
# Instantiate Accessors
2932
self.apps = AppMethods(self.tapis)
3033
self.files = FileMethods(self.tapis)
3134
self.jobs = JobMethods(self.tapis)
35+
self.db = DatabaseAccessor()
3236

3337

3438
# --- AppMethods and FileMethods remain the same ---

‎dapi/db.py‎

Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
import pandas as pd
2+
from .exceptions import DapiException # Use base dapi exception
3+
4+
# --- Database Dependencies Check ---
5+
try:
6+
from sqlalchemy import create_engine, exc, text
7+
from sqlalchemy.orm import sessionmaker
8+
9+
_SQLALCHEMY_FOUND = True
10+
except ImportError:
11+
_SQLALCHEMY_FOUND = False
12+
13+
# Define dummy classes/functions if sqlalchemy is missing,
14+
# so the client can still load but DB access will fail clearly.
15+
class exc: # Dummy class
16+
class SQLAlchemyError(Exception):
17+
pass
18+
19+
def create_engine(*args, **kwargs):
20+
raise ImportError(
21+
"SQLAlchemy not found. Install with 'pip install SQLAlchemy PyMySQL pandas'"
22+
)
23+
24+
def sessionmaker(*args, **kwargs):
25+
raise ImportError(
26+
"SQLAlchemy not found. Install with 'pip install SQLAlchemy PyMySQL pandas'"
27+
)
28+
29+
def text(s):
30+
return s # Return string if no text object
31+
32+
33+
# --- DSDatabase Class ---
34+
class DSDatabase:
35+
"""
36+
A database utility class for connecting to a specific DesignSafe SQL database.
37+
Instantiated by the DSClient with connection details.
38+
"""
39+
40+
def __init__(self, user, password, host, port, db_name, pool_recycle=3600):
41+
"""
42+
Initializes the DSDatabase instance with connection details.
43+
44+
Args:
45+
user (str): Database username.
46+
password (str): Database password.
47+
host (str): Database host address.
48+
port (int): Database port.
49+
db_name (str): Specific database name to connect to.
50+
pool_recycle (int): SQLAlchemy pool recycle time in seconds.
51+
"""
52+
if not _SQLALCHEMY_FOUND:
53+
raise ImportError(
54+
"Database functionality requires SQLAlchemy, PyMySQL, and pandas. Install with 'pip install SQLAlchemy PyMySQL pandas'"
55+
)
56+
57+
self.db_name = db_name
58+
self.connection_string = (
59+
f"mysql+pymysql://{user}:{password}@{host}:{port}/{db_name}"
60+
)
61+
62+
try:
63+
# Setup the database connection
64+
self.engine = create_engine(
65+
self.connection_string,
66+
pool_recycle=pool_recycle,
67+
)
68+
# Test connection immediately (optional but recommended)
69+
with self.engine.connect() as connection:
70+
print(f"Successfully connected to database '{db_name}' on {host}.")
71+
self.Session = sessionmaker(bind=self.engine)
72+
except exc.SQLAlchemyError as e:
73+
# Provide more context in case of connection errors
74+
raise DapiException(
75+
f"Failed to connect to database '{db_name}' at {host}:{port}. Error: {e}"
76+
) from e
77+
except ImportError as e:
78+
# Catch potential PyMySQL import error if create_engine fails due to it
79+
if "pymysql" in str(e).lower():
80+
raise ImportError(
81+
"Database functionality requires PyMySQL. Install with 'pip install PyMySQL'"
82+
) from e
83+
raise # Re-raise other import errors
84+
85+
def read_sql(self, sql: str, output_type: str = "DataFrame") -> Any:
86+
"""
87+
Executes a SQL query and returns the results.
88+
89+
Args:
90+
sql (str): The SQL query string to be executed.
91+
output_type (str, optional): The format for the query results.
92+
Defaults to 'DataFrame'. Possible values are 'DataFrame' or 'dict'.
93+
94+
Returns:
95+
pandas.DataFrame or list[dict]: The result of the SQL query.
96+
97+
Raises:
98+
ValueError: If the SQL query string is empty or output type is invalid.
99+
DapiException: If a database error occurs during query execution.
100+
"""
101+
if not _SQLALCHEMY_FOUND:
102+
raise ImportError(
103+
"Database functionality requires SQLAlchemy, PyMySQL, and pandas. Install with 'pip install SQLAlchemy PyMySQL pandas'"
104+
)
105+
if not sql:
106+
raise ValueError("SQL query string is required")
107+
if output_type not in ["DataFrame", "dict"]:
108+
raise ValueError('Output type must be either "DataFrame" or "dict"')
109+
110+
session = self.Session()
111+
try:
112+
sql_text = text(sql) # Ensure SQL is treated as literal SQL
113+
if output_type == "DataFrame":
114+
# Use pandas directly with the engine for simplicity
115+
return pd.read_sql_query(sql_text, self.engine)
116+
else:
117+
# Execute using session for list of dicts
118+
result = session.execute(sql_text)
119+
# Use .mappings().all() for easy conversion to list of dicts
120+
return result.mappings().all()
121+
except exc.SQLAlchemyError as e:
122+
raise DapiException(
123+
f"Error executing SQL query on database '{self.db_name}'. Error: {e}"
124+
) from e
125+
except Exception as e:
126+
# Catch other potential errors (like pandas issues)
127+
raise DapiException(
128+
f"An unexpected error occurred during SQL query execution: {e}"
129+
) from e
130+
finally:
131+
session.close()
132+
133+
def close(self):
134+
"""Dispose of the SQLAlchemy engine connection pool."""
135+
if hasattr(self, "engine") and self.engine:
136+
print(f"Closing connection pool for database '{self.db_name}'.")
137+
self.engine.dispose()
138+
139+
def __repr__(self):
140+
return f"<DSDatabase(db='{self.db_name}')>"

‎dapi/db/accessor.py‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
from typing import Dict, Optional
2+
from .config import db_config
3+
from .db import DSDatabase
4+
5+
6+
class DatabaseAccessor:
7+
"""
8+
Provides access to different DesignSafe database connections via properties.
9+
Manages lazy initialization of DSDatabase instances and connection pools.
10+
"""
11+
12+
def __init__(self):
13+
self._connections: Dict[str, Optional[DSDatabase]] = {
14+
key: None for key in db_config.keys()
15+
}
16+
print(
17+
"DatabaseAccessor initialized. Connections will be created on first access."
18+
)
19+
20+
def _get_db(self, dbname: str) -> DSDatabase:
21+
"""Gets or creates a DSDatabase instance (with its engine/pool)."""
22+
if dbname not in self._connections:
23+
raise ValueError(
24+
f"Invalid db shorthand '{dbname}'. Allowed: {', '.join(self._connections.keys())}"
25+
)
26+
27+
if self._connections[dbname] is None:
28+
print(f"First access to '{dbname}', initializing DSDatabase...")
29+
try:
30+
self._connections[dbname] = DSDatabase(dbname=dbname)
31+
except Exception as e:
32+
self._connections[dbname] = None
33+
print(f"Error initializing database '{dbname}': {e}")
34+
raise
35+
# Type hint assertion
36+
return self._connections[dbname] # type: ignore
37+
38+
@property
39+
def ngl(self) -> DSDatabase:
40+
"""Access the NGL database connection manager."""
41+
return self._get_db("ngl")
42+
43+
@property
44+
def vp(self) -> DSDatabase:
45+
"""Access the VP database connection manager."""
46+
return self._get_db("vp")
47+
48+
@property
49+
def eq(self) -> DSDatabase:
50+
"""Access the EQ database connection manager."""
51+
return self._get_db("eq")
52+
53+
def close_all(self):
54+
"""Closes all active database engines and their connection pools."""
55+
print("Closing all active database engines/pools...")
56+
closed_count = 0
57+
for dbname, db_instance in self._connections.items():
58+
if db_instance is not None:
59+
try:
60+
# Call the close method on the DSDatabase instance
61+
db_instance.close()
62+
self._connections[
63+
dbname
64+
] = None # Clear instance after closing engine
65+
closed_count += 1
66+
except Exception as e:
67+
print(f"Error closing engine for '{dbname}': {e}")
68+
if closed_count == 0:
69+
print("No active database engines to close.")
70+
else:
71+
print(f"Closed {closed_count} database engine(s).")

‎dapi/db/db.py‎

Lines changed: 48 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -8,33 +8,16 @@
88

99

1010
class DSDatabase:
11-
"""A database utility class for connecting to a DesignSafe SQL database.
12-
13-
This class provides functionality to connect to a MySQL database using
14-
SQLAlchemy and PyMySQL. It supports executing SQL queries and returning
15-
results in different formats.
16-
17-
Attributes:
18-
user (str): Database username, defaults to 'dspublic'.
19-
password (str): Database password, defaults to 'R3ad0nlY'.
20-
host (str): Database host address, defaults to '129.114.52.174'.
21-
port (int): Database port, defaults to 3306.
22-
db (str): Database name, can be 'sjbrande_ngl_db', 'sjbrande_vpdb', or 'post_earthquake_recovery'.
23-
recycle_time (int): Time in seconds to recycle database connections.
24-
engine (Engine): SQLAlchemy engine for database connection.
25-
Session (sessionmaker): SQLAlchemy session maker bound to the engine.
11+
"""
12+
Manages connection and querying for a specific DesignSafe database.
13+
Uses SQLAlchemy engine for connection pooling and session-per-query pattern.
2614
"""
2715

2816
def __init__(self, dbname="ngl"):
29-
"""Initializes the DSDatabase instance with environment variables and creates the database engine.
30-
31-
Args:
32-
dbname (str): Shorthand for the database name. Must be one of 'ngl', 'vp', or 'eq'.
33-
"""
34-
17+
"""Initializes the DSDatabase instance and creates the engine."""
3518
if dbname not in db_config:
3619
raise ValueError(
37-
f"Invalid database shorthand '{dbname}'. Allowed shorthands are: {', '.join(db_config.keys())}"
20+
f"Invalid db shorthand '{dbname}'. Allowed: {', '.join(db_config.keys())}"
3821
)
3922

4023
config = db_config[dbname]
@@ -45,50 +28,71 @@ def __init__(self, dbname="ngl"):
4528
self.host = os.getenv(f"{env_prefix}DB_HOST", "129.114.52.174")
4629
self.port = os.getenv(f"{env_prefix}DB_PORT", 3306)
4730
self.db = config["dbname"]
31+
self.dbname_short = dbname # Store shorthand name for reference
4832

49-
# Setup the database connection
33+
print(
34+
f"Creating SQLAlchemy engine for database '{self.db}' ({self.dbname_short})..."
35+
)
36+
# Setup the database connection engine with pooling
5037
self.engine = create_engine(
5138
f"mysql+pymysql://{self.user}:{self.password}@{self.host}:{self.port}/{self.db}",
52-
pool_recycle=3600, # 1 hour in seconds
39+
pool_recycle=3600, # Recycle connections older than 1 hour
40+
pool_pre_ping=True, # Check connection validity before use
5341
)
42+
# Create a configured "Session" class
5443
self.Session = sessionmaker(bind=self.engine)
44+
print(f"Engine for '{self.dbname_short}' created.")
5545

5646
def read_sql(self, sql, output_type="DataFrame"):
57-
"""Executes a SQL query and returns the results.
58-
59-
Args:
60-
sql (str): The SQL query string to be executed.
61-
output_type (str, optional): The format for the query results. Defaults to 'DataFrame'.
62-
Possible values are 'DataFrame' for a pandas DataFrame, or 'dict' for a list of dictionaries.
63-
64-
Returns:
65-
pandas.DataFrame or list of dict: The result of the SQL query.
47+
"""
48+
Executes a SQL query using a dedicated session and returns the results.
6649
67-
Raises:
68-
ValueError: If the SQL query string is empty or if the output type is not valid.
69-
SQLAlchemyError: If an error occurs during query execution.
50+
Each call obtains a session (and underlying connection from the pool),
51+
executes the query, and closes the session (returning the connection
52+
to the pool).
7053
"""
7154
if not sql:
7255
raise ValueError("SQL query string is required")
73-
7456
if output_type not in ["DataFrame", "dict"]:
7557
raise ValueError('Output type must be either "DataFrame" or "dict"')
7658

59+
# Obtain a new session for this query
7760
session = self.Session()
78-
61+
print(f"Executing query on '{self.dbname_short}'...")
7962
try:
8063
if output_type == "DataFrame":
81-
return pd.read_sql_query(sql, session.bind)
64+
# pandas read_sql_query handles connection/session management implicitly sometimes,
65+
# but using the session explicitly ensures consistency.
66+
# Pass the engine bound to the session.
67+
return pd.read_sql_query(
68+
sql, session.bind.connect()
69+
) # Get connection from engine
8270
else:
83-
# Convert SQL string to a text object
8471
sql_text = text(sql)
72+
# Execute within the session context
8573
result = session.execute(sql_text)
86-
return [dict(row) for row in result]
74+
# Fetch results before closing session
75+
data = [
76+
dict(row._mapping) for row in result
77+
] # Use ._mapping for modern SQLAlchemy
78+
return data
8779
except exc.SQLAlchemyError as e:
88-
raise Exception(f"SQLAlchemyError: {e}")
80+
print(f"SQLAlchemyError executing query on '{self.dbname_short}': {e}")
81+
raise # Re-raise the exception
82+
except Exception as e:
83+
print(f"Unexpected error executing query on '{self.dbname_short}': {e}")
84+
raise
8985
finally:
86+
# Ensure the session is closed, returning the connection to the pool
9087
session.close()
88+
# print(f"Session for '{self.dbname_short}' query closed.") # Can be noisy
9189

9290
def close(self):
93-
"""Close the database connection."""
94-
self.engine.dispose()
91+
"""Dispose of the engine and its connection pool for this database."""
92+
if self.engine:
93+
print(f"Disposing engine and closing pool for '{self.dbname_short}'...")
94+
self.engine.dispose()
95+
self.engine = None # Mark as disposed
96+
print(f"Engine for '{self.dbname_short}' disposed.")
97+
else:
98+
print(f"Engine for '{self.dbname_short}' already disposed.")

‎dapi/db_config.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
# Mapping of shorthand names to actual database names and environment prefixes
2+
db_config = {
3+
"ngl": {"dbname": "sjbrande_ngl_db", "env_prefix": "NGL_"},
4+
"vp": {"dbname": "sjbrande_vpdb", "env_prefix": "VP_"},
5+
"eq": {"dbname": "post_earthquake_recovery", "env_prefix": "EQ_"},
6+
}

0 commit comments

Comments
 (0)