chore: Move function out of receivers

src/sentry/receivers/core.py

@@ -1,12 +1,8 @@
-import logging
-from functools import wraps
-
 from click import echo
 from django.conf import settings
 from django.contrib.auth.models import AnonymousUser
 from django.db import connections, router, transaction
 from django.db.models.signals import post_save
-from django.db.utils import OperationalError, ProgrammingError
 
 from sentry.hybridcloud.models.outbox import outbox_context
 from sentry.loader.dynamic_sdk_options import get_default_loader_data
@@ -19,6 +15,7 @@
 from sentry.signals import post_upgrade, project_created
 from sentry.silo.base import SiloMode, region_silo_function
 from sentry.users.services.user.service import user_service
+from sentry.utils.db import handle_db_failure
 from sentry.utils.env import in_test_environment
 from sentry.utils.settings import is_self_hosted
 
@@ -30,22 +27,6 @@
 DEFAULT_SENTRY_PROJECT_ID = 1
 
 
-def handle_db_failure(func, using=None, wrap_in_transaction=True):
- @wraps(func)
- def wrapped(*args, **kwargs):
- try:
- if wrap_in_transaction:
- with transaction.atomic(using or router.db_for_write(Organization)):
- return func(*args, **kwargs)
- else:
- return func(*args, **kwargs)
- except (ProgrammingError, OperationalError):
- logging.exception("Failed processing signal %s", func.__name__)
- return
-
- return wrapped
-
-
 def create_default_projects(**kwds):
  if not (in_test_environment() or is_self_hosted() or settings.DEBUG):
  # No op in production SaaS environments.
@@ -151,20 +132,20 @@ def freeze_option_epoch_for_project(instance, created, app=None, **kwargs):
 # Anything that relies on default objects that may not exist with default
 # fields should be wrapped in handle_db_failure
 post_upgrade.connect(
- handle_db_failure(create_default_projects, wrap_in_transaction=False),
+ handle_db_failure(create_default_projects, model=Organization, wrap_in_transaction=False),
  dispatch_uid="create_default_project",
  weak=False,
  sender=SiloMode.MONOLITH,
 )
 
 post_save.connect(
- handle_db_failure(freeze_option_epoch_for_project),
+ handle_db_failure(freeze_option_epoch_for_project, model=Organization),
  sender=Project,
  dispatch_uid="freeze_option_epoch_for_project",
  weak=False,
 )
 post_save.connect(
- handle_db_failure(create_keys_for_project),
+ handle_db_failure(create_keys_for_project, model=Organization),
  sender=Project,
  dispatch_uid="create_keys_for_project",
  weak=False,

src/sentry/utils/db.py

@@ -1,11 +1,30 @@
+import logging
 from collections.abc import Sequence
 from contextlib import ExitStack
+from functools import wraps
 
 import sentry_sdk
-from django.db import DEFAULT_DB_ALIAS, connections, transaction
+from django.db import DEFAULT_DB_ALIAS, connections, router, transaction
+from django.db.utils import OperationalError, ProgrammingError
 from sentry_sdk.integrations import Integration
 
 
+def handle_db_failure(func, model, wrap_in_transaction: bool = True):
+ @wraps(func)
+ def wrapped(*args, **kwargs):
+ try:
+ if wrap_in_transaction:
+ with transaction.atomic(router.db_for_write(model)):
+ return func(*args, **kwargs)
+ else:
+ return func(*args, **kwargs)
+ except (ProgrammingError, OperationalError):
+ logging.exception("Failed processing signal %s", func.__name__)
+ return
+
+ return wrapped
+
+
 def atomic_transaction(
  using: str | Sequence[str], savepoint: bool = True
 ) -> transaction.Atomic | ExitStack: