Merge pull request #273 from lsst-sqre/tickets/DM-45281

DM-45281: Allow create_database_engine to take a Url

changelog.d/20240717_220712_rra_DM_45281_queue.md

@@ -0,0 +1,3 @@
+### New features
+
+- `safir.database.create_database_engine` now accepts the database URL as a Pydantic `Url` as well as a `str`.

pyproject.toml

@@ -29,6 +29,7 @@ dependencies = [
     "gidgethub<6",
     "httpx>=0.20.0,<1",
     "pydantic>2,<3",
+    "pydantic-core",
     "starlette<1",
     # 23.3.0 excluded due to https://github.com/hynek/structlog/issues/584
     "structlog>=21.2.0,!=23.3.0",

src/safir/database/_connection.py

@@ -6,6 +6,7 @@
 from urllib.parse import quote, urlparse
 
 from pydantic import SecretStr
+from pydantic_core import Url
 from sqlalchemy.exc import OperationalError
 from sqlalchemy.ext.asyncio import (
     AsyncEngine,
@@ -22,7 +23,9 @@
 ]
 
 
-def _build_database_url(url: str, password: str | SecretStr | None) -> str:
+def _build_database_url(
+    url: str | Url, password: str | SecretStr | None
+) -> str:
     """Build the authenticated URL for the database.
 
     The database scheme is forced to ``postgresql+asyncpg`` if it is
@@ -46,6 +49,8 @@ def _build_database_url(url: str, password: str | SecretStr | None) -> str:
         Raised if a password was provided but the connection URL has no
         username.
     """
+    if not isinstance(url, str):
+        url = str(url)
     parsed_url = urlparse(url)
     if parsed_url.scheme == "postgresql":
         parsed_url = parsed_url._replace(scheme="postgresql+asyncpg")
@@ -67,7 +72,7 @@ def _build_database_url(url: str, password: str | SecretStr | None) -> str:
 
 
 def create_database_engine(
-    url: str,
+    url: str | Url,
     password: str | SecretStr | None,
     *,
     isolation_level: str | None = None,

tests/database_test.py

@@ -9,6 +9,7 @@
 import pytest
 import structlog
 from pydantic import BaseModel, SecretStr
+from pydantic_core import Url
 from sqlalchemy import Column, MetaData, String, Table
 from sqlalchemy.exc import OperationalError, ProgrammingError
 from sqlalchemy.future import select
@@ -83,6 +84,12 @@ def test_build_database_url(database_url: str) -> None:
     )
     assert url == "postgresql+asyncpg://foo:[email protected]:5433/foo"
 
+    pydantic_url = Url.build(
+        scheme="postgresql", username="user", host="localhost", path="foo"
+    )
+    url = _build_database_url(pydantic_url, "password")
+    assert url == "postgresql+asyncpg://user:password@localhost/foo"
+
     # Test that the username and password are quoted properly.
     url = _build_database_url(
         "postgresql://foo%[email protected]:4444/foo",