diff --git a/lib/stac-api/runtime/src/stac_api/handler.py b/lib/stac-api/runtime/src/stac_api/handler.py index db520ba..b5d851c 100644 --- a/lib/stac-api/runtime/src/stac_api/handler.py +++ b/lib/stac-api/runtime/src/stac_api/handler.py @@ -20,19 +20,22 @@ logger = logging.getLogger(__name__) -logger.info("fetching pgstac secret") -secret = get_secret_dict(secret_arn_env_var="PGSTAC_SECRET_ARN") -postgres_settings = PostgresSettings( - pghost=secret["host"], - pgdatabase=secret["dbname"], - pguser=secret["username"], - pgpassword=secret["password"], - pgport=int(secret["port"]), -) - _connection_initialized = False +def _build_postgres_settings() -> PostgresSettings: + """Fetch credentials from Secrets Manager and build PostgresSettings.""" + logger.info("fetching pgstac secret") + secret = get_secret_dict(secret_arn_env_var="PGSTAC_SECRET_ARN") + return PostgresSettings( + pghost=secret["host"], + pgdatabase=secret["dbname"], + pguser=secret["username"], + pgpassword=secret["password"], + pgport=int(secret["port"]), + ) + + @register_before_snapshot def on_snapshot(): """ @@ -90,6 +93,7 @@ def on_snap_restore(): app.state.writepool = None # Create fresh connection pool + postgres_settings = _build_postgres_settings() loop.run_until_complete( connect_to_db( app, @@ -111,6 +115,7 @@ def on_snap_restore(): async def startup_event(): """Connect to database on startup.""" logger.info("Setting up DB connection...") + postgres_settings = _build_postgres_settings() await connect_to_db( app, postgres_settings=postgres_settings, @@ -145,7 +150,7 @@ async def shutdown_event(): loop.run_until_complete( connect_to_db( app, - postgres_settings=postgres_settings, + postgres_settings=_build_postgres_settings(), add_write_connection_pool=with_transactions, ) ) diff --git a/lib/tipg-api/runtime/src/tipg_api/handler.py b/lib/tipg-api/runtime/src/tipg_api/handler.py index 7eefe93..ec6e39d 100644 --- a/lib/tipg-api/runtime/src/tipg_api/handler.py +++ b/lib/tipg-api/runtime/src/tipg_api/handler.py @@ -17,20 +17,24 @@ ) from utils import get_secret_dict -secret = get_secret_dict(secret_arn_env_var="PGSTAC_SECRET_ARN") -postgres_settings = PostgresSettings( - postgres_host=secret["host"], - postgres_dbname=secret["dbname"], - postgres_user=secret["username"], - postgres_pass=secret["password"], - postgres_port=int(secret["port"]), -) db_settings = DatabaseSettings() custom_sql_settings = CustomSQLSettings() _connection_initialized = False +def _build_postgres_settings() -> PostgresSettings: + """Fetch credentials from Secrets Manager and build PostgresSettings.""" + secret = get_secret_dict(secret_arn_env_var="PGSTAC_SECRET_ARN") + return PostgresSettings( + postgres_host=secret["host"], + postgres_dbname=secret["dbname"], + postgres_user=secret["username"], + postgres_pass=secret["password"], + postgres_port=int(secret["port"]), + ) + + @register_before_snapshot def on_snapshot(): """ @@ -74,6 +78,7 @@ def on_snap_restore(): app.state.pool = None # Create fresh connection pool + postgres_settings = _build_postgres_settings() loop.run_until_complete( connect_to_db( app, @@ -103,6 +108,7 @@ def on_snap_restore(): @app.on_event("startup") async def startup_event() -> None: """Connect to database on startup.""" + postgres_settings = _build_postgres_settings() await connect_to_db( app, schemas=db_settings.schemas, diff --git a/lib/titiler-pgstac-api/runtime/src/titiler_pgstac_api/handler.py b/lib/titiler-pgstac-api/runtime/src/titiler_pgstac_api/handler.py index 54ca7ad..a599ae8 100644 --- a/lib/titiler-pgstac-api/runtime/src/titiler_pgstac_api/handler.py +++ b/lib/titiler-pgstac-api/runtime/src/titiler_pgstac_api/handler.py @@ -12,18 +12,21 @@ from titiler.pgstac.settings import PostgresSettings from utils import get_secret_dict -secret = get_secret_dict(secret_arn_env_var="PGSTAC_SECRET_ARN") -postgres_settings = PostgresSettings( - pghost=secret["host"], - pgdatabase=secret["dbname"], - pguser=secret["username"], - pgpassword=secret["password"], - pgport=int(secret["port"]), -) - _connection_initialized = False +def _build_postgres_settings() -> PostgresSettings: + """Fetch credentials from Secrets Manager and build PostgresSettings.""" + secret = get_secret_dict(secret_arn_env_var="PGSTAC_SECRET_ARN") + return PostgresSettings( + pghost=secret["host"], + pgdatabase=secret["dbname"], + pguser=secret["username"], + pgpassword=secret["password"], + pgport=int(secret["port"]), + ) + + @register_before_snapshot def on_snapshot(): """ @@ -67,7 +70,7 @@ def on_snap_restore(): app.state.dbpool = None # Create fresh connection pool - loop.run_until_complete(connect_to_db(app, settings=postgres_settings)) + loop.run_until_complete(connect_to_db(app, settings=_build_postgres_settings())) _connection_initialized = True @@ -81,7 +84,7 @@ def on_snap_restore(): @app.on_event("startup") async def startup_event() -> None: """Connect to database on startup.""" - await connect_to_db(app, settings=postgres_settings) + await connect_to_db(app, settings=_build_postgres_settings()) handler = Mangum(app, lifespan="off")