diff --git a/titiler/pgstac/db.py b/titiler/pgstac/db.py index 8566d91..dd7933e 100644 --- a/titiler/pgstac/db.py +++ b/titiler/pgstac/db.py @@ -1,6 +1,6 @@ """Database connection handling.""" -from typing import Optional +from typing import Any, Dict, Optional from fastapi import FastAPI from psycopg_pool import ConnectionPool @@ -9,12 +9,20 @@ async def connect_to_db( - app: FastAPI, settings: Optional[PostgresSettings] = None + app: FastAPI, + settings: Optional[PostgresSettings] = None, + pool_kwargs: Optional[Dict[str, Any]] = None, ) -> None: """Connect to Database.""" if not settings: settings = PostgresSettings() + pool_kwargs = ( + pool_kwargs + if pool_kwargs is not None + else {"options": "-c search_path=pgstac,public -c application_name=pgstac"} + ) + app.state.dbpool = ConnectionPool( conninfo=str(settings.database_url), min_size=settings.db_min_conn_size, @@ -22,7 +30,7 @@ async def connect_to_db( max_waiting=settings.db_max_queries, max_idle=settings.db_max_idle, num_workers=settings.db_num_workers, - kwargs={"options": "-c search_path=pgstac,public -c application_name=pgstac"}, + kwargs=pool_kwargs, ) # Make sure the pool is ready