-
Notifications
You must be signed in to change notification settings - Fork 16
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[Issue 2665] Implement db connection pools (#2828)
## Summary Maybe Fixes #2665 ### Time to review: __5 mins__ ## Changes proposed > What was added, updated, or removed in this PR. Added connection pools to `/analytics`, replacing single instance db connections. Implementation follows pattern in `/api/src/adapters/db/clients/`. ## Context for reviewers > Testing instructions, background context, more in-depth details of the implementation, and anything else you'd like to call out or ask reviewers. Explain how the changes were verified. This is latest step in a series of attempts to resolve failed db connections from /analytics step functions. See ticket history and comments for details. ## Additional information > Screenshots, GIF demos, code examples or output to help show the changes working as expected.
- Loading branch information
1 parent
f17ff75
commit 2e5d6a8
Showing
4 changed files
with
114 additions
and
58 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,48 +1,77 @@ | ||
# pylint: disable=invalid-name, line-too-long | ||
"""Get a connection to the database using a SQLAlchemy engine object.""" | ||
|
||
from typing import Any, cast | ||
|
||
import boto3 | ||
from sqlalchemy import Engine, create_engine | ||
import psycopg | ||
from sqlalchemy import Connection, Engine, create_engine, pool | ||
|
||
from config import DBSettings, get_db_settings | ||
|
||
# The variables used in the connection url are pulled from local.env | ||
# and configured in the DBSettings class found in config.py | ||
|
||
|
||
def get_db() -> Engine: | ||
""" | ||
Get a connection to the database using a SQLAlchemy engine object. | ||
class PostgresDbClient: | ||
"""An implementation of of a Postgres db client.""" | ||
|
||
def __init__(self, config: DBSettings | None = None) -> None: | ||
"""Construct a class instance.""" | ||
if not config: | ||
config = get_db_settings() | ||
self._engine = self._configure_engine(config) | ||
|
||
def _configure_engine(self, config: DBSettings) -> Engine: | ||
"""Configure db engine to use short-lived IAM tokens for access.""" | ||
|
||
# inspired by /api/src/adapters/db/clients/postgres_client.py | ||
def get_conn() -> psycopg.Connection: | ||
"""Get a psycopg connection.""" | ||
return psycopg.connect(**get_connection_parameters(config)) | ||
|
||
conn_pool = pool.QueuePool(cast(Any, get_conn), max_overflow=5, pool_size=10) | ||
|
||
return create_engine( | ||
"postgresql+psycopg://", | ||
pool=conn_pool, | ||
hide_parameters=True, | ||
) | ||
|
||
def connect(self) -> Connection: | ||
"""Get a new database connection object.""" | ||
return self._engine.connect() | ||
|
||
This function retrieves the database connection URL from the configuration | ||
and creates a SQLAlchemy engine object. | ||
def engine(self) -> Engine: | ||
"""Get reference to db engine.""" | ||
return self._engine | ||
|
||
Yields | ||
------ | ||
sqlalchemy.engine.Engine | ||
A SQLAlchemy engine object representing the connection to the database. | ||
""" | ||
db = get_db_settings() | ||
# inspired by simpler-grants-gov/blob/main/api/src/adapters/db/clients/postgres_client.py | ||
token = db.password if db.local_env is True else generate_iam_auth_token(db) | ||
url = f"postgresql+psycopg://{db.user}:{token}@{db.db_host}:{db.port}/{db.name}?sslmode={db.ssl_mode}" | ||
|
||
return create_engine( | ||
url, | ||
pool_pre_ping=True, | ||
hide_parameters=True, | ||
def get_connection_parameters(config: DBSettings) -> dict[str, Any]: | ||
"""Get parameters for db connection.""" | ||
token = ( | ||
config.password if config.local_env is True else generate_iam_auth_token(config) | ||
) | ||
return { | ||
"host": config.db_host, | ||
"dbname": config.name, | ||
"user": config.user, | ||
"password": token, | ||
"port": config.port, | ||
"connect_timeout": 20, | ||
"sslmode": config.ssl_mode, | ||
} | ||
|
||
|
||
def generate_iam_auth_token(settings: DBSettings) -> str: | ||
def generate_iam_auth_token(config: DBSettings) -> str: | ||
"""Generate IAM auth token.""" | ||
if settings.aws_region is None: | ||
if config.aws_region is None: | ||
msg = "AWS region needs to be configured for DB IAM auth" | ||
raise ValueError(msg) | ||
client = boto3.client("rds", region_name=settings.aws_region) | ||
client = boto3.client("rds", region_name=config.aws_region) | ||
return client.generate_db_auth_token( | ||
DBHostname=settings.db_host, | ||
Port=settings.port, | ||
DBUsername=settings.user, | ||
Region=settings.aws_region, | ||
DBHostname=config.db_host, | ||
Port=config.port, | ||
DBUsername=config.user, | ||
Region=config.aws_region, | ||
) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters