refactor read/write connection pool

Author: vincentsarago

.github/workflows/cicd.yaml

@@ -93,8 +93,7 @@ jobs:
           POSTGRES_USER: username
           POSTGRES_PASS: password
           POSTGRES_DBNAME: postgis
-          POSTGRES_HOST_READER: localhost
-          POSTGRES_HOST_WRITER: localhost
+          POSTGRES_HOST: localhost
           POSTGRES_PORT: 5432
           PGUSER: username
           PGPASSWORD: password

CHANGES.md

@@ -4,11 +4,25 @@
 
 ### Changed
 
+- rename `POSTGRES_HOST_READER` to `POSTGRES_HOST` in config **breaking change**
+- rename `reader_connection_string` to `connection_string` in `PostgresSettings` class **breaking change**
 - add `ENABLE_TRANSACTIONS_EXTENSIONS` env variable to enable `transaction` extensions
 - disable transaction and bulk_transactions extensions by default **breaking change**
 - update `stac-fastapi-*` version requirements to `>=5.2,<6.0`
 - add pgstac health-check in `/_mgmt/health`
 
+### Added
+
+- add `write_connection_pool` option in `stac_fastapi.pgstac.db.connect_to_db` function
+- add `write_postgres_settings` option in `stac_fastapi.pgstac.db.connect_to_db` function to set specific settings for the `writer` DB connection pool
+
+### removed
+
+- `stac_fastapi.pgstac.db.DB` class
+- `POSTGRES_HOST_WRITER` in config
+- `writer_connection_string` in `PostgresSettings` class
+- `testing_connection_string` in `PostgresSettings` class
+
 ## [5.0.2] - 2025-04-07
 
 ### Fixed

docker-compose.yml

@@ -11,8 +11,7 @@ services:
       - POSTGRES_USER=username
       - POSTGRES_PASS=password
       - POSTGRES_DBNAME=postgis
-      - POSTGRES_HOST_READER=database
-      - POSTGRES_HOST_WRITER=database
+      - POSTGRES_HOST=database
       - POSTGRES_PORT=5432
       - WEB_CONCURRENCY=10
       - VSI_CACHE=TRUE

stac_fastapi/pgstac/app.py

@@ -95,7 +95,12 @@
 
 application_extensions = []
 
-if os.environ.get("ENABLE_TRANSACTIONS_EXTENSIONS", "").lower() in ["yes", "true", "1"]:
+with_transactions = os.environ.get("ENABLE_TRANSACTIONS_EXTENSIONS", "").lower() in [
+    "yes",
+    "true",
+    "1",
+]
+if with_transactions:
     application_extensions.append(
         TransactionExtension(
             client=TransactionsClient(),
@@ -150,7 +155,7 @@
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     """FastAPI Lifespan."""
-    await connect_to_db(app)
+    await connect_to_db(app, add_write_connection_pool=with_transactions)
     yield
     await close_db_connection(app)
 

stac_fastapi/pgstac/config.py

@@ -54,8 +54,7 @@ class PostgresSettings(BaseSettings):
     Attributes:
         postgres_user: postgres username.
         postgres_pass: postgres password.
-        postgres_host_reader: hostname for the reader connection.
-        postgres_host_writer: hostname for the writer connection.
+        postgres_host: hostname for the connection.
         postgres_port: database port.
         postgres_dbname: database name.
         use_api_hydrate: perform hydration of stac items within stac-fastapi.
@@ -64,8 +63,7 @@ class PostgresSettings(BaseSettings):
 
     postgres_user: str
     postgres_pass: str
-    postgres_host_reader: str
-    postgres_host_writer: str
+    postgres_host: str
     postgres_port: int
     postgres_dbname: str
 
@@ -79,19 +77,9 @@ class PostgresSettings(BaseSettings):
     model_config = {"env_file": ".env", "extra": "ignore"}
 
     @property
-    def reader_connection_string(self):
+    def connection_string(self):
         """Create reader psql connection string."""
-        return f"postgresql://{self.postgres_user}:{quote(self.postgres_pass)}@{self.postgres_host_reader}:{self.postgres_port}/{self.postgres_dbname}"
-
-    @property
-    def writer_connection_string(self):
-        """Create writer psql connection string."""
-        return f"postgresql://{self.postgres_user}:{quote(self.postgres_pass)}@{self.postgres_host_writer}:{self.postgres_port}/{self.postgres_dbname}"
-
-    @property
-    def testing_connection_string(self):
-        """Create testing psql connection string."""
-        return f"postgresql://{self.postgres_user}:{quote(self.postgres_pass)}@{self.postgres_host_writer}:{self.postgres_port}/pgstactestdb"
+        return f"postgresql://{self.postgres_user}:{quote(self.postgres_pass)}@{self.postgres_host}:{self.postgres_port}/{self.postgres_dbname}"
 
 
 class Settings(ApiSettings):

stac_fastapi/pgstac/db.py

@@ -13,11 +13,10 @@
     Union,
 )
 
-import attr
 import orjson
-from asyncpg import Connection, exceptions
+from asyncpg import Connection, Pool, exceptions
 from buildpg import V, asyncpg, render
-from fastapi import FastAPI, Request
+from fastapi import FastAPI, HTTPException, Request
 from stac_fastapi.types.errors import (
     ConflictError,
     DatabaseError,
@@ -47,33 +46,46 @@ async def con_init(conn):
 ConnectionGetter = Callable[[Request, Literal["r", "w"]], AsyncIterator[Connection]]
 
 
+async def _create_pool(settings: PostgresSettings) -> Pool:
+    """Create a connection pool."""
+    return await asyncpg.create_pool(
+        settings.connection_string,
+        min_size=settings.db_min_conn_size,
+        max_size=settings.db_max_conn_size,
+        max_queries=settings.db_max_queries,
+        max_inactive_connection_lifetime=settings.db_max_inactive_conn_lifetime,
+        init=con_init,
+        server_settings=settings.server_settings.model_dump(),
+    )
+
+
 async def connect_to_db(
     app: FastAPI,
     get_conn: Optional[ConnectionGetter] = None,
     postgres_settings: Optional[PostgresSettings] = None,
+    add_write_connection_pool: bool = False,
+    write_postgres_settings: Optional[PostgresSettings] = None,
 ) -> None:
     """Create connection pools & connection retriever on application."""
-    app_settings = app.state.settings
-
     if not postgres_settings:
         postgres_settings = PostgresSettings()
 
-    if app_settings.testing:
-        readpool = writepool = postgres_settings.testing_connection_string
-    else:
-        readpool = postgres_settings.reader_connection_string
-        writepool = postgres_settings.writer_connection_string
+    app.state.readpool = await _create_pool(postgres_settings)
+
+    if add_write_connection_pool:
+        if not write_postgres_settings:
+            write_postgres_settings = PostgresSettings()
+
+        app.state.writepool = await _create_pool(write_postgres_settings)
 
-    db = DB()
-    app.state.readpool = await db.create_pool(readpool, postgres_settings)
-    app.state.writepool = await db.create_pool(writepool, postgres_settings)
     app.state.get_connection = get_conn if get_conn else get_connection
 
 
 async def close_db_connection(app: FastAPI) -> None:
     """Close connection."""
     await app.state.readpool.close()
-    await app.state.writepool.close()
+    if pool := getattr(app.state, "writepool", None):
+        await pool.close()
 
 
 @asynccontextmanager
@@ -82,7 +94,15 @@ async def get_connection(
     readwrite: Literal["r", "w"] = "r",
 ) -> AsyncIterator[Connection]:
     """Retrieve connection from database conection pool."""
-    pool = request.app.state.writepool if readwrite == "w" else request.app.state.readpool
+    pool = request.app.state.readpool
+    if readwrite == "w":
+        pool = getattr(request.app.state, "writepool", None)
+        if not pool:
+            raise HTTPException(
+                status_code=500,
+                detail="Could not find connection pool for write operations",
+            )
+
     with translate_pgstac_errors():
         async with pool.acquire() as conn:
             yield conn
@@ -131,25 +151,3 @@ def translate_pgstac_errors() -> Generator[None, None, None]:
         raise DatabaseError from e
     except exceptions.ForeignKeyViolationError as e:
         raise ForeignKeyError from e
-
-
-@attr.s
-class DB:
-    """DB class that can be used with context manager."""
-
-    connection_string = attr.ib(default=None)
-    _pool = attr.ib(default=None)
-    _connection = attr.ib(default=None)
-
-    async def create_pool(self, connection_string: str, settings):
-        """Create a connection pool."""
-        pool = await asyncpg.create_pool(
-            connection_string,
-            min_size=settings.db_min_conn_size,
-            max_size=settings.db_max_conn_size,
-            max_queries=settings.db_max_queries,
-            max_inactive_connection_lifetime=settings.db_max_inactive_conn_lifetime,
-            init=con_init,
-            server_settings=settings.server_settings.model_dump(),
-        )
-        return pool

stac_fastapi/pgstac/extensions/filter.py

@@ -25,9 +25,7 @@ async def get_queryables(
         under OGC CQL but it is allowed by the STAC API Filter Extension
         https://github.com/radiantearth/stac-api-spec/tree/master/fragments/filter#queryables
         """
-        pool = request.app.state.readpool
-
-        async with pool.acquire() as conn:
+        async with request.app.state.get_connection(request, "r") as conn:
             q, p = render(
                 """
                     SELECT * FROM get_queryables(:collection::text);

tests/api/test_api.py

@@ -742,8 +742,7 @@ async def get_collection(
     postgres_settings = PostgresSettings(
         postgres_user=database.user,
         postgres_pass=database.password,
-        postgres_host_reader=database.host,
-        postgres_host_writer=database.host,
+        postgres_host=database.host,
         postgres_port=database.port,
         postgres_dbname=database.dbname,
     )
@@ -770,7 +769,12 @@ async def get_collection(
         collections_get_request_model=collection_search_extension.GET,
     )
     app = api.app
-    await connect_to_db(app, postgres_settings=postgres_settings)
+    await connect_to_db(
+        app,
+        postgres_settings=postgres_settings,
+        add_write_connection_pool=True,
+        write_postgres_settings=postgres_settings,
+    )
     try:
         async with AsyncClient(transport=ASGITransport(app=app)) as client:
             response = await client.post(
@@ -812,8 +816,7 @@ async def test_no_extension(
     postgres_settings = PostgresSettings(
         postgres_user=database.user,
         postgres_pass=database.password,
-        postgres_host_reader=database.host,
-        postgres_host_writer=database.host,
+        postgres_host=database.host,
         postgres_port=database.port,
         postgres_dbname=database.dbname,
     )
@@ -826,7 +829,12 @@ async def test_no_extension(
         search_post_request_model=post_request_model,
     )
     app = api.app
-    await connect_to_db(app, postgres_settings=postgres_settings)
+    await connect_to_db(
+        app,
+        postgres_settings=postgres_settings,
+        add_write_connection_pool=True,
+        write_postgres_settings=postgres_settings,
+    )
     try:
         async with AsyncClient(transport=ASGITransport(app=app)) as client:
             landing = await client.get("http://test/")

tests/clients/test_postgres.py

@@ -529,8 +529,7 @@ async def test_db_setup_works_with_env_vars(api_client, database, monkeypatch):
     """Test that the application starts successfully if the POSTGRES_* environment variables are set"""
     monkeypatch.setenv("POSTGRES_USER", database.user)
     monkeypatch.setenv("POSTGRES_PASS", database.password)
-    monkeypatch.setenv("POSTGRES_HOST_READER", database.host)
-    monkeypatch.setenv("POSTGRES_HOST_WRITER", database.host)
+    monkeypatch.setenv("POSTGRES_HOST", database.host)
     monkeypatch.setenv("POSTGRES_PORT", str(database.port))
     monkeypatch.setenv("POSTGRES_DBNAME", database.dbname)
 
@@ -567,8 +566,7 @@ async def app(self, api_client, database):
         postgres_settings = PostgresSettings(
             postgres_user=database.user,
             postgres_pass=database.password,
-            postgres_host_reader=database.host,
-            postgres_host_writer=database.host,
+            postgres_host=database.host,
             postgres_port=database.port,
             postgres_dbname=database.dbname,
         )

tests/conftest.py

@@ -202,15 +202,19 @@ async def app(api_client, database):
     postgres_settings = PostgresSettings(
         postgres_user=database.user,
         postgres_pass=database.password,
-        postgres_host_reader=database.host,
-        postgres_host_writer=database.host,
+        postgres_host=database.host,
         postgres_port=database.port,
         postgres_dbname=database.dbname,
     )
     logger.info("Creating app Fixture")
     time.time()
     app = api_client.app
-    await connect_to_db(app, postgres_settings=postgres_settings)
+    await connect_to_db(
+        app,
+        postgres_settings=postgres_settings,
+        add_write_connection_pool=True,
+        write_postgres_settings=postgres_settings,
+    )
 
     yield app
 
@@ -306,14 +310,18 @@ async def app_no_ext(database):
     postgres_settings = PostgresSettings(
         postgres_user=database.user,
         postgres_pass=database.password,
-        postgres_host_reader=database.host,
-        postgres_host_writer=database.host,
+        postgres_host=database.host,
         postgres_port=database.port,
         postgres_dbname=database.dbname,
     )
     logger.info("Creating app Fixture")
     time.time()
-    await connect_to_db(api_client_no_ext.app, postgres_settings=postgres_settings)
+    await connect_to_db(
+        api_client_no_ext.app,
+        postgres_settings=postgres_settings,
+        add_write_connection_pool=True,
+        write_postgres_settings=postgres_settings,
+    )
     yield api_client_no_ext.app
     await close_db_connection(api_client_no_ext.app)
 
@@ -343,14 +351,17 @@ async def app_no_transaction(database):
     postgres_settings = PostgresSettings(
         postgres_user=database.user,
         postgres_pass=database.password,
-        postgres_host_reader=database.host,
-        postgres_host_writer=database.host,
+        postgres_host=database.host,
         postgres_port=database.port,
         postgres_dbname=database.dbname,
     )
     logger.info("Creating app Fixture")
     time.time()
-    await connect_to_db(api.app, postgres_settings=postgres_settings)
+    await connect_to_db(
+        api.app,
+        postgres_settings=postgres_settings,
+        add_write_connection_pool=False,
+    )
     yield api.app
     await close_db_connection(api.app)
 
@@ -371,8 +382,7 @@ async def default_app(database, monkeypatch):
     """Test default stac-fastapi-pgstac application."""
     monkeypatch.setenv("POSTGRES_USER", database.user)
     monkeypatch.setenv("POSTGRES_PASS", database.password)
-    monkeypatch.setenv("POSTGRES_HOST_READER", database.host)
-    monkeypatch.setenv("POSTGRES_HOST_WRITER", database.host)
+    monkeypatch.setenv("POSTGRES_HOST", database.host)
     monkeypatch.setenv("POSTGRES_PORT", str(database.port))
     monkeypatch.setenv("POSTGRES_DBNAME", database.dbname)
     monkeypatch.delenv("ENABLED_EXTENSIONS", raising=False)
@@ -383,7 +393,7 @@ async def default_app(database, monkeypatch):
 
     from stac_fastapi.pgstac.app import app
 
-    await connect_to_db(app)
+    await connect_to_db(app, add_write_connection_pool=True)
     yield app
     await close_db_connection(app)