feat(aurora): invalidate connection on readonly transaction error - DIA-51197 (#79)

Author: arththebird

fastapi_sqla/asyncio_support.py

@@ -12,7 +12,8 @@
 from sqlalchemy.orm.session import sessionmaker
 from sqlalchemy.sql import Select, func, select
 
-from fastapi_sqla.sqla import Base, Page, T, aws_rds_iam_support, new_engine
+from fastapi_sqla import aws_aurora_support, aws_rds_iam_support
+from fastapi_sqla.sqla import Base, Page, T, new_engine
 
 logger = structlog.get_logger(__name__)
 _ASYNC_SESSION_KEY = "fastapi_sqla_async_session"
@@ -31,6 +32,7 @@ def new_async_engine():
 async def startup():
     engine = new_async_engine()
     aws_rds_iam_support.setup(engine.sync_engine)
+    aws_aurora_support.setup(engine.sync_engine)
 
     # Fail early:
     try:

fastapi_sqla/aws_aurora_support.py

@@ -0,0 +1,33 @@
+from pydantic import BaseSettings
+from sqlalchemy import event
+from sqlalchemy.engine import Engine
+from sqlalchemy.engine.interfaces import ExceptionContext
+
+# Taken from
+# https://www.postgresql.org/docs/current/errcodes-appendix.html#ERRCODES-TABLE
+READONLY_ERROR_CODE = "25006"
+
+
+def setup(engine: Engine):
+    config = Config()
+
+    if not config.aws_aurora_enabled:
+        return
+
+    event.listen(engine, "handle_error", disconnect_on_readonly_error)
+
+
+def disconnect_on_readonly_error(context: ExceptionContext):
+    if context.is_disconnect:
+        return
+
+    error_code = getattr(context.original_exception, "pgcode", None)
+    if error_code == READONLY_ERROR_CODE:
+        context.is_disconnect = True  # type: ignore
+
+
+class Config(BaseSettings):
+    aws_aurora_enabled: bool = False
+
+    class Config:
+        env_prefix = "fastapi_sqla_"

fastapi_sqla/sqla.py

@@ -20,7 +20,7 @@
 from sqlalchemy.orm.session import sessionmaker
 from sqlalchemy.sql import Select, func, select
 
-from fastapi_sqla import aws_rds_iam_support
+from fastapi_sqla import aws_aurora_support, aws_rds_iam_support
 
 try:
     from sqlalchemy.orm import declarative_base
@@ -50,6 +50,7 @@ def is_async_dialect(engine):
 def startup():
     engine = new_engine()
     aws_rds_iam_support.setup(engine.engine)
+    aws_aurora_support.setup(engine.engine)
 
     # Fail early:
     try:

tests/test_aws_aurora_support.py

@@ -0,0 +1,68 @@
+from unittest.mock import Mock
+
+from pytest import mark, raises
+from sqlalchemy import text
+
+
+@mark.sqlalchemy("1.4")
+@mark.dont_patch_engines
+def test_sync_disconnects_on_readonly_error(monkeypatch):
+    from fastapi_sqla.sqla import _Session, startup
+
+    monkeypatch.setenv("fastapi_sqla_aws_aurora_enabled", "true")
+
+    startup()
+
+    connection = _Session().connection(execution_options={"postgresql_readonly": True})
+    with raises(Exception):
+        connection.execute(text("CREATE TABLE fail(id integer)"))
+
+    assert connection.invalidated
+
+
+@mark.sqlalchemy("1.4")
+@mark.require_asyncpg
+@mark.dont_patch_engines
+async def test_async_disconnects_on_readonly_error(monkeypatch, async_sqlalchemy_url):
+    from fastapi_sqla.asyncio_support import _AsyncSession, startup
+
+    monkeypatch.setenv("fastapi_sqla_aws_aurora_enabled", "true")
+    monkeypatch.setenv("async_sqlalchemy_url", async_sqlalchemy_url)
+
+    await startup()
+
+    connection = await _AsyncSession().connection(
+        execution_options={"postgresql_readonly": True}
+    )
+    with raises(Exception):
+        await connection.execute(text("CREATE TABLE fail(id integer)"))
+
+    assert connection.invalidated
+
+
+@mark.parametrize(
+    "pgcode, should_disconnect", [("25006", True), ("00000", False), (None, False)]
+)
+def test_readonly_error_codes(pgcode, should_disconnect):
+    from fastapi_sqla.aws_aurora_support import disconnect_on_readonly_error
+
+    exception = Mock()
+    exception.pgcode = pgcode
+    context = Mock()
+    context.original_exception = exception
+    context.is_disconnect = False
+
+    disconnect_on_readonly_error(context)
+
+    assert context.is_disconnect == should_disconnect
+
+
+def test_already_disconnected():
+    from fastapi_sqla.aws_aurora_support import disconnect_on_readonly_error
+
+    context = Mock()
+    context.is_disconnect = True
+
+    disconnect_on_readonly_error(context)
+
+    assert context.is_disconnect is True