Add tooling to make customization of DB Connection possible

[alukach]

Description:

This PR updates the way in which the pgSTAC extension connects with its underlying database. This was done to support the possibility of users running custom connection setup/teardown steps for each request. This allows us to do something like the following when setting up a pgSTAC-based STAC API:

diff --git a/stac_fastapi/pgstac/stac_fastapi/pgstac/app.py b/stac_fastapi/pgstac/stac_fastapi/pgstac/app.py
index 8a33424..1033f11 100644
--- a/stac_fastapi/pgstac/stac_fastapi/pgstac/app.py
+++ b/stac_fastapi/pgstac/stac_fastapi/pgstac/app.py
@@ -6,8 +6,14 @@ If the variable is not set, enables all extensions.
 """
 
 import os
+from contextlib import asynccontextmanager
+from typing import Literal
 
+from asyncpg import exceptions, Connection
+from buildpg import render
+from fastapi import HTTPException, Request
 from fastapi.responses import ORJSONResponse
+from fastapi.security import HTTPBasic
 
 from stac_fastapi.api.app import StacApi
 from stac_fastapi.api.models import create_get_request_model, create_post_request_model
@@ -22,7 +28,7 @@ from stac_fastapi.extensions.core import (
 from stac_fastapi.extensions.third_party import BulkTransactionExtension
 from stac_fastapi.pgstac.config import Settings
 from stac_fastapi.pgstac.core import CoreCrudClient
-from stac_fastapi.pgstac.db import close_db_connection, connect_to_db
+from stac_fastapi.pgstac.db import close_db_connection, connect_to_db, get_connection
 from stac_fastapi.pgstac.extensions import QueryExtension
 from stac_fastapi.pgstac.extensions.filter import FiltersClient
 from stac_fastapi.pgstac.transactions import BulkTransactionsClient, TransactionsClient
@@ -65,10 +71,57 @@ api = StacApi(
 app = api.app
 
 
+auth = HTTPBasic()
+
+
+async def set_db_user(conn: Connection, request: Request) -> None:
+    # Require & parse basic auth
+    credentials = await auth(request)
+
+    # Validate credentials...
+    q, p = render(
+        f"""
+        SELECT authenticate_user(:username::text, :password::text)
+        """,
+        username=credentials.username,
+        password=credentials.password,
+    )
+    try:
+        await conn.execute(q, *p)
+    except exceptions.RaiseError as e:
+        raise HTTPException(status_code=401, detail="Bad credentials") from e
+
+    # Set user in DB...
+    q, p = render(
+        f"""
+        SELECT set_config('api.username', :username::text, false)
+        """,
+        username=credentials.username,
+    )
+    await conn.execute(q, *p)
+
+
+@asynccontextmanager
+async def custom_get_connection(
+    request: Request,
+    readwrite: Literal["r", "w"],
+):
+    """An example of customizing the connection getter"""
+    async with get_connection(request, readwrite) as conn:
+        if readwrite == "w":
+            await set_db_user(conn, request)
+
+        # Simple demo to validate we now know who is authenticated
+        current_user = await conn.fetchval("SELECT current_setting('api.username', true)")
+        print(f"Running DB command as {current_user!r}")
+        
+        yield conn
+
+
 @app.on_event("startup")
 async def startup_event():
     """Connect to database on startup."""
-    await connect_to_db(app)
+    await connect_to_db(app, custom_get_connection)
 
 
 @app.on_event("shutdown")

PR Checklist:

• Code is formatted and linted (run pre-commit run --all-files)
• Tests pass (run make test)
• Documentation has been updated to reflect changes, if applicable, and docs build successfully (run make docs)
• Changes are added to the CHANGELOG.

[bitner]

The other thing I'll note that this can be used for is for platforms that do not allow setting search_path on config. The SET SEARCH_PATH TO pgstac, public; could be done using a connection init function.

[gadomski]

Just a note that the CI failures are due to non-smart configuration of the GHCR updates (by me), and should be fixed by #535 (which will close #534).

[alukach]

Alright @bitner & others, I believe this is now good to go. Testing is failing due to the issue described in stac-utils/pgstac#177, but that should be fixed when we upgrade to pgSTAC 0.7.5 (via fix from stac-utils/pgstac#178).

I added one simple test to validate that the customized DB connection can be used to run logic with API requests:

stac-fastapi/stac_fastapi/pgstac/tests/clients/test_postgres.py

Lines 182 to 212 in befe6ea

	@asynccontextmanager
	async def custom_get_connection(
	request: Request,
	readwrite: Literal["r", "w"],
	):
	"""An example of customizing the connection getter"""
	async with get_connection(request, readwrite) as conn:
	await conn.execute("SELECT set_config('api.test', 'added-config', false)")
	yield conn
	class TestDbConnect:
	@pytest.fixture
	async def app(self, api_client):
	"""
	app fixture override to setup app with a customized db connection getter
	"""
	logger.debug("Customizing app setup")
	await connect_to_db(api_client.app, custom_get_connection)
	yield api_client.app
	await close_db_connection(api_client.app)
	async def test_db_setup(self, api_client, app_client):
	@api_client.app.get(f"{api_client.router.prefix}/db-test")
	async def example_view(request: Request):
	async with request.app.state.get_connection(request, "r") as conn:
	return await conn.fetchval("SELECT current_setting('api.test', true)")
	response = await app_client.get("/db-test")
	assert response.status_code == 200
	assert response.json() == "added-config"

[bitner]

Looks good! Thank you @alukach. @gadomski I just want to make sure that you see this too as far as it relates to the backend separation PR.

[gadomski]

Sounds good, thanks for the heads up. I'll merge this as is, and then open an issue to backport it to https://github.com/stac-utils/stac-fastapi-pgstac/.

[gadomski]

stac-utils/stac-fastapi-pgstac#19