fix(api): update database routing logic in MainRouter (#9080)

api/src/backend/tasks/jobs/integrations.py

@@ -5,7 +5,7 @@ from celery.utils.log import get_task_logger
 from config.django.base import DJANGO_FINDINGS_BATCH_SIZE
 from tasks.utils import batched
 
-from api.db_router import READ_REPLICA_ALIAS
+from api.db_router import READ_REPLICA_ALIAS, MainRouter
 from api.db_utils import rls_transaction
 from api.models import Finding, Integration, Provider
 from api.utils import initialize_prowler_integration, initialize_prowler_provider
@@ -179,7 +179,7 @@ def get_security_hub_client_from_integration(
         if the connection was successful and the SecurityHub client or connection object.
     """
     # Get the provider associated with this integration
-    with rls_transaction(tenant_id):
+    with rls_transaction(tenant_id, using=READ_REPLICA_ALIAS):
         provider_relationship = integration.integrationproviderrelationship_set.first()
         if not provider_relationship:
             return Connection(
@@ -208,7 +208,7 @@ def get_security_hub_client_from_integration(
             regions_status[region] = region in connection.enabled_regions
 
         # Save regions information in the integration configuration
-        with rls_transaction(tenant_id):
+        with rls_transaction(tenant_id, using=MainRouter.default_db):
             integration.configuration["regions"] = regions_status
             integration.save()
 
@@ -223,7 +223,7 @@ def get_security_hub_client_from_integration(
         return True, security_hub
     else:
         # Reset regions information if connection fails
-        with rls_transaction(tenant_id):
+        with rls_transaction(tenant_id, using=MainRouter.default_db):
             integration.configuration["regions"] = {}
             integration.save()
 
@@ -334,8 +334,11 @@ def upload_security_hub_integration(
                                         f"Security Hub connection failed for integration {integration.id}: "
                                         f"{security_hub.error}"
                                     )
-                                    integration.connected = False
-                                    integration.save()
+                                    with rls_transaction(
+                                        tenant_id, using=MainRouter.default_db
+                                    ):
+                                        integration.connected = False
+                                        integration.save()
                                     break  # Skip this integration
 
                                 security_hub_client = security_hub

api/src/backend/tasks/tests/test_integrations.py

@@ -9,6 +9,7 @@ from tasks.jobs.integrations import (
     upload_security_hub_integration,
 )
 
+from api.db_router import READ_REPLICA_ALIAS, MainRouter
 from api.models import Integration
 from api.utils import prowler_integration_connection_test
 from prowler.providers.aws.lib.security_hub.security_hub import SecurityHubConnection
@@ -880,7 +881,8 @@ class TestSecurityHubIntegrationUploads:
         # Verify RLS transaction was used correctly
         # Should be called twice: once for getting provider info, once for resetting regions
         assert mock_rls.call_count == 2
-        mock_rls.assert_any_call(tenant_id)
+        mock_rls.assert_any_call(tenant_id, using=READ_REPLICA_ALIAS)
+        mock_rls.assert_any_call(tenant_id, using=MainRouter.default_db)
 
         # Verify test_connection was called with integration credentials (not provider's)
         mock_test_connection.assert_called_once_with(