feat: add get_by_id_for_update method to LoginSessionRepository and implement... (c55b38f4) · Commits · OOP / code / portal

src/portal_service/adapters/database/repos/login_sessions.py

+10 −0

Changes for src/portal_service/adapters/database/repos/login_sessions.py: 10 added lines, 0 removed lines.
Original line number	Diff line number	Diff line
		@@ -19,6 +19,16 @@ class SqlLoginSessionRepository(LoginSessionRepository):
		row = await self._session.scalar(stmt)
		return LoginSessionMapper.to_domain(row) if row is not None else None

		async def get_by_id_for_update(self, session_id_hash: str) -> LoginSession \| None:
		stmt = (
		select(LoginSessionRow)
		.where(LoginSessionRow.id == session_id_hash)
		.with_for_update()
		.execution_options(populate_existing=True)
		)
		row = await self._session.scalar(stmt)
		return LoginSessionMapper.to_domain(row) if row is not None else None

		async def save(self, login_session: LoginSession) -> LoginSession:
		merged = await self._session.merge(LoginSessionMapper.to_row(login_session))
		await self._session.flush()

src/portal_service/application/session_service.py

+25 −10

Changes for src/portal_service/application/session_service.py: 25 added lines, 10 removed lines.
Original line number	Diff line number	Diff line
		@@ -39,6 +39,11 @@ class ActiveSession:
		_REFRESH_LEEWAY = timedelta(seconds=30)


		def _is_fresh(session: LoginSession, now: datetime) -> bool:
		"""The access token is good for long enough to use without refreshing."""
		return session.access_token_expires_at - _REFRESH_LEEWAY > now


		class SessionService:
		def __init__(
		self,
		@@ -100,34 +105,44 @@ class SessionService:
		We drop the session when that happens, so the next request is a clean logout.
		"""
		now = datetime.now(UTC)
		if active.session.access_token_expires_at - _REFRESH_LEEWAY > now:
		if _is_fresh(active.session, now):
		return active
		if active.tokens.refresh_token is None:

		# A page load sends several requests at once, and they all find the token
		# expired. Locking the row lets one of them refresh; the others wait here and
		# then use its result. Without this each would present the same refresh token,
		# which Keycloak refuses once rotation is on, and the refusal logs the user out.
		current = await self._sessions.get_by_id_for_update(active.session.id)
		if current is None or current.expires_at <= now:
		raise InvalidSessionError("Session ended while waiting to refresh")
		tokens = decrypt_tokens(current.encrypted_tokens, active.credentials)
		if _is_fresh(current, now):
		# Another request refreshed it while we waited.
		return ActiveSession(session=current, tokens=tokens, credentials=active.credentials)
		if tokens.refresh_token is None:
		raise InvalidSessionError("Session has no refresh token")

		try:
		issued = await self._identity_provider.refresh(active.tokens.refresh_token)
		issued = await self._identity_provider.refresh(tokens.refresh_token)
		except InvalidGrantError as exc:
		await self._sessions.delete_by_id(active.session.id)
		await self._sessions.delete_by_id(current.id)
		raise InvalidSessionError("Keycloak no longer accepts this session") from exc

		# Keycloak may return only some of the tokens; keep what it did not replace.
		tokens = SessionTokens(
		access_token=issued.tokens.access_token,
		refresh_token=issued.tokens.refresh_token or active.tokens.refresh_token,
		id_token=issued.tokens.id_token or active.tokens.id_token,
		refresh_token=issued.tokens.refresh_token or tokens.refresh_token,
		id_token=issued.tokens.id_token or tokens.id_token,
		)
		access_token_expires_at = now + timedelta(seconds=issued.expires_in)
		# Same key as before: it lives only in the cookie the browser still holds.
		encrypted = encrypt_tokens(tokens, active.credentials)
		if not await self._sessions.update_tokens(
		active.session.id, encrypted, access_token_expires_at
		):
		if not await self._sessions.update_tokens(current.id, encrypted, access_token_expires_at):
		# Logged out while we were refreshing: the session stays ended.
		raise InvalidSessionError("Session ended during refresh")

		# The session's own deadline is unchanged by a refresh.
		row = active.session.model_copy(
		row = current.model_copy(
		update={
		"encrypted_tokens": encrypted,
		"access_token_expires_at": access_token_expires_at,

src/portal_service/domain/ports/database/login_sessions.py

+7 −0

Changes for src/portal_service/domain/ports/database/login_sessions.py: 7 added lines, 0 removed lines.
Original line number	Diff line number	Diff line
		@@ -16,6 +16,13 @@ class LoginSessionRepository(ABC):
		async def get_by_id(self, session_id_hash: str) -> LoginSession \| None:
		pass

		@abstractmethod
		async def get_by_id_for_update(self, session_id_hash: str) -> LoginSession \| None:
		"""
		Like get_by_id, but reads the row as it is now and locks it until the
		transaction ends, so concurrent refreshes of one session run one at a time.
		"""

		@abstractmethod
		async def save(self, login_session: LoginSession) -> LoginSession:
		pass

tests/fakes.py

+4 −0

Changes for tests/fakes.py: 4 added lines, 0 removed lines.
Original line number	Diff line number	Diff line
		@@ -33,6 +33,10 @@ class FakeLoginSessionRepository(LoginSessionRepository):
		async def get_by_id(self, session_id_hash: str) -> LoginSession \| None:
		return self.rows.get(session_id_hash)

		async def get_by_id_for_update(self, session_id_hash: str) -> LoginSession \| None:
		# One request at a time in these tests: nothing to lock.
		return self.rows.get(session_id_hash)

		async def save(self, login_session: LoginSession) -> LoginSession:
		self.rows[login_session.id] = login_session
		return login_session

tests/integration/test_concurrent_refresh.py

0 → 100644

+107 −0

Changes for tests/integration/test_concurrent_refresh.py: 107 added lines, 0 removed lines.
Original line number	Diff line number	Diff line
		"""
		Concurrent refreshes of one session, against a real Postgres.

		A page load sends several requests at once, and right after the access token
		expires they all want to refresh it. Only the row lock makes them take turns, and
		an in-memory fake cannot show that, so these run on separate database connections.
		"""

		import asyncio
		from collections.abc import AsyncIterator
		from datetime import UTC, datetime, timedelta

		import pytest
		from sqlalchemy import delete
		from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine

		from portal_service.adapters.database.repos import SqlLoginSessionRepository
		from portal_service.adapters.database.sql import LoginSessionRow, get_metadata
		from portal_service.application.session_crypto import SessionCredentials, encrypt_tokens
		from portal_service.application.session_service import ActiveSession, SessionService
		from portal_service.domain.errors import InvalidGrantError
		from portal_service.domain.models import IssuedTokens, LoginSession, SessionTokens
		from tests.fakes import FakeIdentityProvider

		pytestmark = pytest.mark.integration


		class RotatingIdentityProvider(FakeIdentityProvider):
		"""Keycloak with refresh-token rotation on: each refresh token works once."""

		def __init__(self) -> None:
		super().__init__()
		self._used: set[str] = set()

		async def refresh(self, refresh_token: str) -> IssuedTokens:
		self.refreshed.append(refresh_token)
		if refresh_token in self._used:
		raise InvalidGrantError("Maximum allowed refresh token reuse exceeded")
		self._used.add(refresh_token)
		# Slow enough that the other request is certainly waiting meanwhile.
		await asyncio.sleep(0.3)
		n = len(self._used)
		return IssuedTokens(
		tokens=SessionTokens(access_token=f"access-{n}", refresh_token=f"refresh-{n}"),
		expires_in=300,
		)


		@pytest.fixture
		async def session_maker(postgres_url: str) -> AsyncIterator[async_sessionmaker[AsyncSession]]:
		"""Committing sessions on separate connections, as concurrent requests have."""
		engine = create_async_engine(postgres_url)
		async with engine.begin() as connection:
		await connection.run_sync(get_metadata().create_all)
		try:
		yield async_sessionmaker(bind=engine, expire_on_commit=False)
		finally:
		async with engine.begin() as connection:
		await connection.execute(delete(LoginSessionRow))
		await engine.dispose()


		async def _expired_session(maker: async_sessionmaker[AsyncSession]) -> str:
		"""Stores a session whose access token has expired; returns its cookie."""
		credentials = SessionCredentials.generate()
		now = datetime.now(UTC)
		async with maker() as db:
		await SqlLoginSessionRepository(db).save(
		LoginSession(
		id=credentials.session_id_hash,
		subject="user-123",
		access_token_expires_at=now - timedelta(seconds=10),
		encrypted_tokens=encrypt_tokens(
		SessionTokens(access_token="access-0", refresh_token="refresh-0"),
		credentials,
		),
		expires_at=now + timedelta(hours=1),
		)
		)
		await db.commit()
		return credentials.to_cookie()


		async def _request(
		maker: async_sessionmaker[AsyncSession], idp: FakeIdentityProvider, cookie: str
		) -> ActiveSession:
		"""What one request does: load, refresh if needed, commit."""
		async with maker() as db:
		service = SessionService(SqlLoginSessionRepository(db), idp, timedelta(hours=1), "/")
		active = await service.ensure_fresh(await service.load_without_refreshing(cookie))
		await db.commit()
		return active


		async def test_concurrent_requests_refresh_once_and_all_succeed(
		session_maker: async_sessionmaker[AsyncSession],
		) -> None:
		cookie = await _expired_session(session_maker)
		idp = RotatingIdentityProvider()

		results = await asyncio.gather(*(_request(session_maker, idp, cookie) for _ in range(3)))

		# One trip to Keycloak; the others waited and used its tokens.
		assert idp.refreshed == ["refresh-0"]
		assert {r.tokens.access_token for r in results} == {"access-1"}
		# And the session survived for the next request.
		assert (await _request(session_maker, idp, cookie)).tokens.access_token == "access-1"