Merge pull request #156 from stac-utils/backport/pool-kwargs-option

backport change from 1.2.2

CHANGES.md

@@ -1,5 +1,9 @@
 # Release Notes
 
+## 0.8.3 (2024-02-21)
+
+* enable passing `ConnectionPool` kwargs option in `titiler.pgstac.db.connect_to_db` function (author @smohiudd, #155) [backported from 1.2.2]
+
 ## 0.8.2 (2024-01-23)
 
 * update rio-tiler version to `>6.3.0` (defined in `titiler>=0.17`)

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,22 +9,34 @@
 
 
 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,
         max_size=settings.db_max_conn_size,
         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
+    # ref: https://www.psycopg.org/psycopg3/docs/advanced/pool.html#pool-startup-check
+    app.state.dbpool.wait()
+
 
 async def close_db_connection(app: FastAPI) -> None:
     """Close Pool."""