fix(api): make social signup transactional (#12245)

api/changelog.d/social-signup-transaction.fixed.md

@@ -0,0 +1 @@
+Social signups create users and authentication records in one database transaction, preventing incomplete accounts when provisioning fails

api/src/backend/api/adapters.py

@@ -1,7 +1,7 @@
 from allauth.account.models import EmailAddress
 from allauth.core.exceptions import ImmediateHttpResponse
 from allauth.socialaccount.adapter import DefaultSocialAccountAdapter
-from api.db_router import MainRouter
+from api.db_router import MainRouter, write_db_alias
 from api.db_utils import rls_transaction
 from api.models import (
     Membership,
@@ -107,7 +107,10 @@ class ProwlerSocialAccountAdapter(DefaultSocialAccountAdapter):
         and is about to be saved to the DB for the first time.
         """
         with transaction.atomic(using=MainRouter.admin_db):
-            user = super().save_user(request, sociallogin, form)
+            # Allauth saves the user without an explicit alias. Route that save
+            # through admin so every signup record shares this transaction.
+            with write_db_alias(MainRouter.admin_db):
+                user = super().save_user(request, sociallogin, form)
             provider = sociallogin.provider.id
             extra = sociallogin.account.extra_data
 

api/src/backend/api/db_router.py

@@ -1,3 +1,4 @@
+from contextlib import contextmanager
 from contextvars import ContextVar
 
 from django.conf import settings
@@ -5,6 +6,7 @@ from django.conf import settings
 ALLOWED_APPS = ("django", "socialaccount", "account", "authtoken", "silk")
 
 _read_db_alias = ContextVar("read_db_alias", default=None)
+_write_db_alias = ContextVar("write_db_alias", default=None)
 
 
 def set_read_db_alias(alias: str | None):
@@ -22,6 +24,30 @@ def reset_read_db_alias(token) -> None:
         _read_db_alias.reset(token)
 
 
+def set_write_db_alias(alias: str | None):
+    if not alias:
+        return None
+    return _write_db_alias.set(alias)
+
+
+def get_write_db_alias() -> str | None:
+    return _write_db_alias.get()
+
+
+def reset_write_db_alias(token) -> None:
+    if token is not None:
+        _write_db_alias.reset(token)
+
+
+@contextmanager
+def write_db_alias(alias: str | None):
+    token = set_write_db_alias(alias)
+    try:
+        yield
+    finally:
+        reset_write_db_alias(token)
+
+
 class MainRouter:
     default_db = "default"
     admin_db = "admin"
@@ -43,6 +69,9 @@ class MainRouter:
         model_table_name = model._meta.db_table
         if any(model_table_name.startswith(f"{app}_") for app in ALLOWED_APPS):
             return self.admin_db
+        write_alias = get_write_db_alias()
+        if write_alias:
+            return write_alias
         return None
 
     def allow_migrate(self, db, app_label, model_name=None, **hints):  # noqa: F841

api/src/backend/api/tests/test_adapters.py

@@ -10,10 +10,12 @@ from allauth.socialaccount import app_settings as socialaccount_app_settings
 from allauth.socialaccount.internal.flows.login import complete_login
 from allauth.socialaccount.models import SocialAccount, SocialLogin
 from api.adapters import ProwlerSocialAccountAdapter
-from api.db_router import MainRouter
+from api.db_router import MainRouter, get_write_db_alias
 from api.models import Invitation, Membership, SAMLConfiguration, Tenant
 from django.contrib.auth import get_user_model
 from django.core import mail
+from django.db import connections
+from django.db import router as django_router
 
 User = get_user_model()
 
@@ -382,6 +384,65 @@ class TestProwlerSocialAccountAdapter:
             role=Membership.RoleChoices.MEMBER,
         ).exists()
 
+    def test_save_user_routes_initial_allauth_write_to_admin_and_resets_on_error(
+        self, rf
+    ):
+        adapter = ProwlerSocialAccountAdapter()
+        request = rf.get("/")
+        request.session = {}
+        sociallogin = _oauth_sociallogin(
+            User(name="Frank", email="frank-routing@example.com")
+        )
+
+        def fail_after_checking_write_route(*_args, **_kwargs):
+            assert (
+                MainRouter().db_for_write(User, instance=sociallogin.user)
+                == MainRouter.admin_db
+            )
+            raise RuntimeError("Stop after checking the write route.")
+
+        with (
+            patch("api.adapters.super") as mock_super,
+            patch("api.adapters.transaction.atomic"),
+            patch.object(MainRouter, "admin_db", "admin"),
+            pytest.raises(RuntimeError, match="Stop after checking the write route"),
+        ):
+            mock_super.return_value.save_user.side_effect = (
+                fail_after_checking_write_route
+            )
+            adapter.save_user(request, sociallogin)
+
+        assert get_write_db_alias() is None
+
+    def test_save_user_rolls_back_all_signup_records_on_downstream_error(self, rf):
+        adapter = ProwlerSocialAccountAdapter()
+        request = rf.post("/")
+        request.session = {}
+        email = "frank-rollback@example.com"
+        sociallogin = _real_oauth_sociallogin(
+            User(name="Frank", email=email),
+            uid="frank-rollback-google-account",
+        )
+        tenants_before = Tenant.objects.count()
+
+        with (
+            patch(
+                "api.adapters.rls_transaction",
+                side_effect=RuntimeError("Simulated downstream failure."),
+            ),
+            pytest.raises(RuntimeError, match="Simulated downstream failure"),
+        ):
+            adapter.save_user(request, sociallogin)
+
+        assert not User.objects.filter(email=email).exists()
+        assert not SocialAccount.objects.filter(
+            provider="google",
+            uid="frank-rollback-google-account",
+        ).exists()
+        assert not EmailAddress.objects.filter(email=email).exists()
+        assert Tenant.objects.count() == tenants_before
+        assert get_write_db_alias() is None
+
     def test_save_user_saml_sets_session_flag(self, rf):
         adapter = ProwlerSocialAccountAdapter()
         request = rf.get("/")
@@ -402,3 +463,104 @@ class TestProwlerSocialAccountAdapter:
                     mock_super.return_value.save_user.return_value = mock_user
                     adapter.save_user(request, sociallogin)
                     assert request.session["saml_user_created"] == "123"
+
+
+@pytest.mark.requires_test_admin_alias
+@pytest.mark.django_db(transaction=True, databases=["default", "admin"])
+class TestProwlerSocialAccountAdapterMultiDatabase:
+    @staticmethod
+    def _production_router():
+        return patch.object(django_router, "routers", [MainRouter()])
+
+    def test_save_user_rolls_back_across_production_database_aliases(self, rf):
+        adapter = ProwlerSocialAccountAdapter()
+        request = rf.post("/")
+        request.session = {}
+        email = "frank-multidb-rollback@example.com"
+        sociallogin = _real_oauth_sociallogin(
+            User(name="Frank", email=email),
+            uid="frank-multidb-rollback-google-account",
+        )
+        tenants_before = Tenant.objects.using("admin").count()
+
+        assert connections["default"] is not connections["admin"]
+        assert (
+            connections["default"].settings_dict["NAME"]
+            == connections["admin"].settings_dict["NAME"]
+        )
+
+        def fail_after_allauth_save(*_args, **_kwargs):
+            assert sociallogin.user._state.db == MainRouter.admin_db
+            assert connections["default"].get_autocommit()
+            assert not connections["admin"].get_autocommit()
+            raise RuntimeError("Simulated downstream failure.")
+
+        with (
+            patch.object(MainRouter, "admin_db", "admin"),
+            self._production_router(),
+            patch("api.adapters.rls_transaction", side_effect=fail_after_allauth_save),
+            pytest.raises(RuntimeError, match="Simulated downstream failure"),
+        ):
+            adapter.save_user(request, sociallogin)
+
+        assert connections["default"].get_autocommit()
+        assert connections["admin"].get_autocommit()
+        assert not User.objects.using("default").filter(email=email).exists()
+        assert not User.objects.using("admin").filter(email=email).exists()
+        assert (
+            not SocialAccount.objects.using("admin")
+            .filter(
+                provider="google",
+                uid="frank-multidb-rollback-google-account",
+            )
+            .exists()
+        )
+        assert not EmailAddress.objects.using("admin").filter(email=email).exists()
+        assert Tenant.objects.using("admin").count() == tenants_before
+        assert get_write_db_alias() is None
+
+    def test_save_user_commits_complete_signup_across_production_aliases(self, rf):
+        adapter = ProwlerSocialAccountAdapter()
+        request = rf.post("/")
+        request.session = {}
+        email = "frank-multidb-success@example.com"
+        sociallogin = _real_oauth_sociallogin(
+            User(name="Frank", email=email),
+            uid="frank-multidb-success-google-account",
+        )
+
+        with (
+            patch.object(MainRouter, "admin_db", "admin"),
+            self._production_router(),
+        ):
+            user = adapter.save_user(request, sociallogin)
+
+        user = User.objects.using("admin").get(id=user.id)
+        assert user.email == email
+        assert (
+            SocialAccount.objects.using("admin")
+            .filter(
+                user_id=user.id,
+                provider="google",
+                uid="frank-multidb-success-google-account",
+            )
+            .exists()
+        )
+        assert (
+            EmailAddress.objects.using("admin")
+            .filter(
+                user_id=user.id,
+                email=email,
+                verified=True,
+            )
+            .exists()
+        )
+        assert (
+            Membership.objects.using("admin")
+            .filter(
+                user_id=user.id,
+                role=Membership.RoleChoices.OWNER,
+            )
+            .exists()
+        )
+        assert get_write_db_alias() is None

api/src/backend/api/tests/test_database.py

@@ -1,7 +1,13 @@
-from unittest.mock import patch
+from unittest.mock import Mock, patch
 
 import pytest
-from api.db_router import MainRouter
+from api.db_router import (
+    MainRouter,
+    get_write_db_alias,
+    reset_write_db_alias,
+    set_write_db_alias,
+    write_db_alias,
+)
 from api.rls import Tenant
 from config.django.base import DATABASE_ROUTERS as PROD_DATABASE_ROUTERS
 from django.conf import settings
@@ -26,6 +32,66 @@ class TestMainDatabaseRouter:
         assert router.allow_migrate_model(MainRouter.admin_db, api_model)
         assert not router.allow_migrate_model("default", api_model)
 
+    def test_scoped_write_alias_routes_api_models(self, router):
+        token = set_write_db_alias(MainRouter.admin_db)
+        try:
+            assert get_write_db_alias() == MainRouter.admin_db
+            assert router.db_for_write(Tenant) == MainRouter.admin_db
+        finally:
+            reset_write_db_alias(token)
+
+        assert get_write_db_alias() is None
+        assert router.db_for_write(Tenant) == "default"
+
+    def test_scoped_write_alias_restores_nested_context(self, router):
+        outer_token = set_write_db_alias("outer")
+        try:
+            assert router.db_for_write(Tenant) == "outer"
+
+            inner_token = set_write_db_alias(MainRouter.admin_db)
+            try:
+                assert router.db_for_write(Tenant) == MainRouter.admin_db
+            finally:
+                reset_write_db_alias(inner_token)
+
+            assert router.db_for_write(Tenant) == "outer"
+        finally:
+            reset_write_db_alias(outer_token)
+
+        assert get_write_db_alias() is None
+        assert router.db_for_write(Tenant) == "default"
+
+    def test_scoped_write_alias_does_not_override_admin_models(self, router):
+        token = set_write_db_alias("other")
+        try:
+            assert (
+                router.db_for_write(MigrationRecorder.Migration) == MainRouter.admin_db
+            )
+        finally:
+            reset_write_db_alias(token)
+
+        assert get_write_db_alias() is None
+
+    def test_write_db_alias_context_manager_resets_after_error(self, router):
+        fail = Mock(side_effect=RuntimeError("Simulated failure"))
+
+        with pytest.raises(RuntimeError, match="Simulated failure"):
+            with write_db_alias(MainRouter.admin_db):
+                assert get_write_db_alias() == MainRouter.admin_db
+                assert router.db_for_write(Tenant) == MainRouter.admin_db
+                fail()
+
+        fail.assert_called_once_with()
+        assert get_write_db_alias() is None
+        assert router.db_for_write(Tenant) == "default"
+
+    def test_write_db_alias_context_manager_ignores_empty_alias(self, router):
+        with write_db_alias(None):
+            assert get_write_db_alias() is None
+            assert router.db_for_write(Tenant) == "default"
+
+        assert get_write_db_alias() is None
+
     def test_router_django_models(self, router):
         assert router.db_for_read(MigrationRecorder.Migration) == MainRouter.admin_db
         assert not router.db_for_read(MigrationRecorder.Migration) == "default"

api/src/backend/conftest.py

@@ -70,6 +70,7 @@ API_JSON_CONTENT_TYPE = "application/vnd.api+json"
 NO_TENANT_HTTP_STATUS = status.HTTP_401_UNAUTHORIZED
 TEST_USER = "dev@prowler.com"
 TEST_PASSWORD = "testing_psswd"
+TEST_ADMIN_ALIAS = "admin"
 TEST_REPLICA_ALIAS = "test_replica"
 
 
@@ -2539,26 +2540,36 @@ def finding_groups_title_variants_fixture(
     return findings
 
 
+def _ensure_mirrored_test_alias(alias: str) -> None:
+    default_database = settings.DATABASES["default"]
+    if alias not in settings.DATABASES:
+        settings.DATABASES[alias] = {
+            **default_database,
+            "TEST": {
+                **default_database.get("TEST", {}),
+                "MIRROR": "default",
+            },
+        }
+    django_connections.databases[alias] = settings.DATABASES[alias]
+
+
 def pytest_collection_modifyitems(items):
     """Ensure test_rbac.py is executed first."""
     items.sort(key=lambda item: 0 if "test_rbac.py" in item.nodeid else 1)
 
+    if any(item.get_closest_marker("requires_test_admin_alias") for item in items):
+        _ensure_mirrored_test_alias(TEST_ADMIN_ALIAS)
+
     if any(item.get_closest_marker("requires_test_replica_alias") for item in items):
-        default_database = settings.DATABASES["default"]
-        if TEST_REPLICA_ALIAS not in settings.DATABASES:
-            settings.DATABASES[TEST_REPLICA_ALIAS] = {
-                **default_database,
-                "TEST": {
-                    **default_database.get("TEST", {}),
-                    "MIRROR": "default",
-                },
-            }
-        django_connections.databases[TEST_REPLICA_ALIAS] = settings.DATABASES[
-            TEST_REPLICA_ALIAS
-        ]
+        _ensure_mirrored_test_alias(TEST_REPLICA_ALIAS)
 
 
 def pytest_configure(config):
+    config.addinivalue_line(
+        "markers",
+        "requires_test_admin_alias: creates a test-only admin alias mirrored "
+        "to default",
+    )
     config.addinivalue_line(
         "markers",
         "requires_test_replica_alias: creates a test-only replica alias mirrored "