From 15c46d268f24107153ce27cb0f22a3f789ddbf48 Mon Sep 17 00:00:00 2001
From: Jo Booth <jo.m.booth@gmail.com>
Date: Fri, 4 Aug 2023 17:23:33 -0400
Subject: [PATCH] fix: LSDV-5289: add logging to transactions FF check for
 debugging purposes (#4603)

fix: LSDV-5289: add logging to FF check for debugging purposes
---
 label_studio/core/utils/db.py   | 31 +++++++++++++++++++++++++++----
 label_studio/projects/models.py | 11 +++++------
 label_studio/tasks/models.py    |  6 +++---
 3 files changed, 35 insertions(+), 13 deletions(-)

diff --git a/label_studio/core/utils/db.py b/label_studio/core/utils/db.py
index e11a618e273..21f19d6e772 100644
--- a/label_studio/core/utils/db.py
+++ b/label_studio/core/utils/db.py
@@ -1,8 +1,14 @@
+from typing import TYPE_CHECKING
+import logging
+
+from core.feature_flags import flag_set
 from django.db import models
+from django.db.models import Subquery
+
+if TYPE_CHECKING:
+    from users.models import User
 
-from django.db.models import (
-    Subquery
-)
+logger = logging.getLogger(__name__)
 
 
 class SQCount(Subquery):
@@ -12,9 +18,26 @@ class SQCount(Subquery):
 
 def fast_first(queryset):
     """Replacement for queryset.first() when you don't need ordering,
-       queryset.first() works slowly in some cases
+    queryset.first() works slowly in some cases
     """
     try:
         return queryset.all()[0]
     except IndexError:
         return None
+
+
+def should_run_bulk_update_in_transaction(organization_created_by_user: "User") -> bool:
+    """Check flag for the given user, log result and user id to info for debugging
+    purposes"""
+
+    bulk_update_should_run_in_transaction = flag_set(
+        "fflag_fix_back_lsdv_5289_run_bulk_updates_in_transactions_short",
+        user=organization_created_by_user,
+        override_system_default=True,
+    )
+
+    logger.info(
+        f"[deadlocks debugging] {bulk_update_should_run_in_transaction=} {organization_created_by_user.id=}"
+    )
+
+    return bulk_update_should_run_in_transaction
diff --git a/label_studio/projects/models.py b/label_studio/projects/models.py
index b4fbd6e665f..644768e5614 100644
--- a/label_studio/projects/models.py
+++ b/label_studio/projects/models.py
@@ -21,6 +21,7 @@
     load_func,
     merge_labels_counters,
 )
+from core.utils.db import should_run_bulk_update_in_transaction
 from core.utils.exceptions import LabelStudioValidationErrorSentryIgnored
 from core.label_config import (
     validate_label_config,
@@ -854,9 +855,8 @@ def _update_tasks_counters(self, queryset, from_scratch=True):
             task.total_predictions = task.new_total_predictions
             objs.append(task)
         with conditional_atomic(
-            predicate=flag_set,
-            predicate_args=['fflag_fix_back_lsdv_5289_run_bulk_updates_in_transactions_short'],
-            predicate_kwargs={'user': self.organization.created_by},
+            predicate=should_run_bulk_update_in_transaction,
+            predicate_args=[self.organization.created_by],
         ):
             bulk_update(objs, update_fields=['total_annotations', 'cancelled_annotations', 'total_predictions'], batch_size=settings.BATCH_SIZE)
         return len(objs)
@@ -874,9 +874,8 @@ def _update_tasks_counters_and_is_labeled(self, task_ids, from_scratch=True):
 
         while (task_ids_slice := task_ids[page_idx * settings.BATCH_SIZE:(page_idx + 1) * settings.BATCH_SIZE]):
             with conditional_atomic(
-                predicate=flag_set,
-                predicate_args=['fflag_fix_back_lsdv_5289_run_bulk_updates_in_transactions_short'],
-                predicate_kwargs={'user': organization_created_by},
+                predicate=should_run_bulk_update_in_transaction,
+                predicate_args=[organization_created_by],
             ):
                 # If counters are updated, is_labeled must be updated as well. Hence, if either fails, we
                 # will roll back. NB: as part of LSDV-5289, we are considering eliminating this transaction
diff --git a/label_studio/tasks/models.py b/label_studio/tasks/models.py
index 1ea8af63a5f..c78c2d21c47 100644
--- a/label_studio/tasks/models.py
+++ b/label_studio/tasks/models.py
@@ -33,6 +33,7 @@
     string_is_url,
     temporary_disconnect_list_signal,
 )
+from core.utils.db import should_run_bulk_update_in_transaction
 from core.utils.params import get_env
 from core.label_config import SINGLE_VALUED_TAGS
 from core.current_request import get_current_request
@@ -897,9 +898,8 @@ def bulk_update_stats_project_tasks(tasks, project=None):
         project = tasks[0].project
 
     with conditional_atomic(
-        predicate=flag_set,
-        predicate_args=['fflag_fix_back_lsdv_5289_run_bulk_updates_in_transactions_short'],
-        predicate_kwargs={'user': project.organization.created_by},
+        predicate=should_run_bulk_update_in_transaction,
+        predicate_args=[project.organization.created_by],
     ):
         use_overlap = project._can_use_overlap()
         maximum_annotations = project.maximum_annotations