Repository: saleor/saleor

Issue:
Bug: `stockBulkUpdate` is slow with large dataset
### What are you trying to achieve?

Run the `stockBulkUpdate` mutation with the following datasets:
- 100 items in the `stocks` input list
- have at least 100K stock records in the DB

At some point, the mutation runs a SELECT query that fetches 100K stock records from DB. A DB transaction that's run in this mutation then takes over 40 seconds to execute. 

### Steps to reproduce the problem

See above.

### What did you expect to happen?

Mutation should be optimized to run in a few seconds at maximum. 

### Logs

DD trace [here](https://app.datadoghq.com/apm/trace/3666491088843823336?colorBy=service&env=production&graphType=flamegraph&shouldShowLegend=true&sort=time&spanID=13138036231616013459&spanViewType=metadata&timeHint=1698372172078.027&view=spans).

### Environment

Saleor version: 3.14.29




Discussion:

Resolving patch:
diff --git a/saleor/graphql/warehouse/bulk_mutations/stock_bulk_update.py b/saleor/graphql/warehouse/bulk_mutations/stock_bulk_update.py
--- a/saleor/graphql/warehouse/bulk_mutations/stock_bulk_update.py
+++ b/saleor/graphql/warehouse/bulk_mutations/stock_bulk_update.py
@@ -83,7 +83,9 @@ class Arguments:
 
     class Meta:
         description = (
-            "Updates stocks for a given variant and warehouse."
+            "Updates stocks for a given variant and warehouse. Variant and warehouse "
+            "selectors have to be the same for all stock inputs. Is not allowed to "
+            "use 'variantId' in one input and 'variantExternalReference' in another."
             + ADDED_IN_313
             + PREVIEW_FEATURE
         )
@@ -190,6 +192,49 @@ def validate_warehouse(
                 errors_count += 1
         return errors_count
 
+    @classmethod
+    def get_selectors(cls, cleaned_inputs_map):
+        return_error = False
+        warehouse_selector = None
+        variant_selector = None
+
+        for stock_input in cleaned_inputs_map.values():
+            if stock_input:
+                warehouse_selector = (
+                    "warehouse_id"
+                    if cleaned_inputs_map[0].get("warehouse_id")
+                    else "warehouse_external_reference"
+                )
+                variant_selector = (
+                    "variant_id"
+                    if cleaned_inputs_map[0].get("variant_id")
+                    else "variant_external_reference"
+                )
+                break
+
+        if not warehouse_selector or not variant_selector:
+            return "warehouse_id", "variant_id"
+
+        # check if all inputs have same selectors
+        if not all(warehouse_selector in s for s in cleaned_inputs_map.values() if s):
+            return_error = True
+
+        if not all(variant_selector in s for s in cleaned_inputs_map.values() if s):
+            return_error = True
+
+        if return_error:
+            message = (
+                "All inputs should use the same selector for "
+                "variant (`variantId` or `variantExternalReference`) and "
+                "warehouse (`warehouseId` or `warehouseExternalReference`)."
+            )
+            raise ValidationError(
+                message=message,
+                code=StockBulkUpdateErrorCode.GRAPHQL_ERROR.value,
+            )
+
+        return warehouse_selector, variant_selector
+
     @classmethod
     def clean_stocks(cls, stocks_input, index_error_map):
         cleaned_inputs_map: dict = {}
@@ -231,36 +276,24 @@ def clean_stocks(cls, stocks_input, index_error_map):
 
     @classmethod
     def _get_stock(
-        cls, variant_id, variant_external_ref, warehouse_id, warehouse_external_ref
+        cls, warehouse_selector, variant_selector, warehouse_value, variant_value
     ):
-        if variant_id and warehouse_id:
-            return (
-                lambda stock: str(stock.warehouse_id) == warehouse_id
-                and str(stock.product_variant_id) == variant_id
-            )
-        elif variant_external_ref and warehouse_id:
-            return (
-                lambda stock: str(stock.warehouse_id) == warehouse_id
-                and stock.variant_external_reference == variant_external_ref
-            )
+        if variant_selector == "variant_id":
+            variant_selector = "product_variant_id"
 
-        elif variant_external_ref and warehouse_external_ref:
-            return (
-                lambda stock: str(stock.warehouse_external_reference)
-                == warehouse_external_ref
-                and stock.variant_external_reference == variant_external_ref
-            )
-        else:
-            return (
-                lambda stock: stock.warehouse_external_reference
-                == warehouse_external_ref
-                and str(stock.product_variant_id) == variant_id
-            )
+        return (
+            lambda stock: str(getattr(stock, warehouse_selector)) == warehouse_value
+            and str(getattr(stock, variant_selector)) == variant_value
+        )
 
     @classmethod
-    def update_stocks(cls, cleaned_inputs_map, index_error_map):
+    def update_stocks(
+        cls, cleaned_inputs_map, warehouse_selector, variant_selector, index_error_map
+    ):
         instances_data_and_errors_list: list = []
-        stocks_list = cls.get_stocks(cleaned_inputs_map)
+        stocks_list = cls.get_stocks(
+            cleaned_inputs_map, warehouse_selector, variant_selector
+        )
 
         for index, cleaned_input in cleaned_inputs_map.items():
             if not cleaned_input:
@@ -269,18 +302,16 @@ def update_stocks(cls, cleaned_inputs_map, index_error_map):
                 )
                 continue
 
-            variant_id = cleaned_input.get("variant_id")
-            warehouse_id = cleaned_input.get("warehouse_id")
-            variant_external_ref = cleaned_input.get("variant_external_reference")
-            warehouse_external_ref = cleaned_input.get("warehouse_external_reference")
+            warehouse_value = cleaned_input.get(warehouse_selector)
+            variant_value = cleaned_input.get(variant_selector)
 
             filter_stock = list(
                 filter(
                     cls._get_stock(
-                        variant_id,
-                        variant_external_ref,
-                        warehouse_id,
-                        warehouse_external_ref,
+                        warehouse_selector,
+                        variant_selector,
+                        warehouse_value,
+                        variant_value,
                     ),
                     stocks_list,
                 )
@@ -304,39 +335,40 @@ def update_stocks(cls, cleaned_inputs_map, index_error_map):
         return instances_data_and_errors_list
 
     @classmethod
-    def get_stocks(cls, cleaned_inputs_map: dict) -> list[models.Stock]:
-        lookup = Q()
-        for stocks_input in cleaned_inputs_map.values():
-            if not stocks_input:
-                continue
+    def get_stocks(
+        cls, cleaned_inputs_map: dict, warehouse_selector: str, variant_selector: str
+    ) -> list[models.Stock]:
+        warehouses = (
+            stock_input[warehouse_selector]
+            for stock_input in cleaned_inputs_map.values()
+            if stock_input
+        )
 
-            single_stock_lookup = Q()
+        variants = (
+            stock_input[variant_selector]
+            for stock_input in cleaned_inputs_map.values()
+            if stock_input
+        )
 
-            if variant_id := stocks_input.get("variant_id"):
-                single_stock_lookup |= Q(product_variant_id=variant_id)
-            else:
-                single_stock_lookup |= Q(
-                    product_variant__external_reference=stocks_input.get(
-                        "variant_external_reference"
-                    )
-                )
+        if not warehouses or not variants:
+            return []
 
-            if warehouse_id := stocks_input.get("warehouse_id"):
-                single_stock_lookup |= Q(warehouse_id=warehouse_id)
-            else:
-                single_stock_lookup |= Q(
-                    warehouse__external_reference=stocks_input.get(
-                        "variant_external_reference"
-                    )
-                )
+        if warehouse_selector == "warehouse_id":
+            warehouse_lookup = Q(warehouse_id__in=warehouses)
+        else:
+            warehouse_lookup = Q(warehouse__external_reference__in=warehouses)
 
-            lookup |= single_stock_lookup
+        if variant_selector == "variant_id":
+            variant_lookup = Q(product_variant_id__in=variants)
+        else:
+            variant_lookup = Q(product_variant__external_reference__in=variants)
 
-        stocks = models.Stock.objects.filter(lookup).annotate(
+        stocks = models.Stock.objects.filter(
+            warehouse_lookup & variant_lookup
+        ).annotate(
             variant_external_reference=F("product_variant__external_reference"),
             warehouse_external_reference=F("warehouse__external_reference"),
         )
-
         return list(stocks)
 
     @classmethod
@@ -385,8 +417,10 @@ def perform_mutation(cls, _root, info, **data):
         index_error_map: dict = defaultdict(list)
         cleaned_inputs_map = cls.clean_stocks(data["stocks"], index_error_map)
 
+        warehouse_selector, variant_selector = cls.get_selectors(cleaned_inputs_map)
+
         instances_data_with_errors_list = cls.update_stocks(
-            cleaned_inputs_map, index_error_map
+            cleaned_inputs_map, warehouse_selector, variant_selector, index_error_map
         )
 
         if any([bool(error) for error in index_error_map.values()]):


Test patch:
diff --git a/saleor/graphql/warehouse/tests/benchmark/test_stock_bulk_update.py b/saleor/graphql/warehouse/tests/benchmark/test_stock_bulk_update.py
--- a/saleor/graphql/warehouse/tests/benchmark/test_stock_bulk_update.py
+++ b/saleor/graphql/warehouse/tests/benchmark/test_stock_bulk_update.py
@@ -1,4 +1,3 @@
-import graphene
 import pytest
 from django.db import connection
 
@@ -51,21 +50,16 @@ def test_stocks_bulk_update_queries_count(
         warehouse=warehouse_no_shipping_zone, product_variant=variant, quantity=4
     )
 
-    variant_1_id = graphene.Node.to_global_id("ProductVariant", variant_1.pk)
-
     stocks = variant.stocks.all()
     stock_1 = stocks[0]
     stock_2 = stocks[1]
 
     new_quantity = 999
 
-    warehouse_1_id = graphene.Node.to_global_id("Warehouse", stock_1.warehouse_id)
-    warehouse_4_id = graphene.Node.to_global_id("Warehouse", stock_4.warehouse_id)
-
     stocks_input = [
         {
-            "variantId": variant_1_id,
-            "warehouseId": warehouse_1_id,
+            "variantExternalReference": variant_1.external_reference,
+            "warehouseExternalReference": stock_1.warehouse.external_reference,
             "quantity": new_quantity,
         }
     ]
@@ -90,7 +84,7 @@ def test_stocks_bulk_update_queries_count(
 
     stocks_input += [
         {
-            "variantId": variant_1_id,
+            "variantExternalReference": variant_1.external_reference,
             "warehouseExternalReference": stock_2.warehouse.external_reference,
             "quantity": new_quantity,
         },
@@ -101,7 +95,7 @@ def test_stocks_bulk_update_queries_count(
         },
         {
             "variantExternalReference": variant_2.external_reference,
-            "warehouseId": warehouse_4_id,
+            "warehouseExternalReference": stock_4.warehouse.external_reference,
             "quantity": new_quantity,
         },
     ]
diff --git a/saleor/graphql/warehouse/tests/test_stock_bulk_update.py b/saleor/graphql/warehouse/tests/test_stock_bulk_update.py
--- a/saleor/graphql/warehouse/tests/test_stock_bulk_update.py
+++ b/saleor/graphql/warehouse/tests/test_stock_bulk_update.py
@@ -8,6 +8,10 @@
 STOCKS_BULK_UPDATE_MUTATION = """
     mutation StockBulkUpdate($stocks: [StockBulkUpdateInput!]!){
         stockBulkUpdate(stocks: $stocks){
+            errors{
+                message
+                code
+            }
             results{
                 errors {
                     field
@@ -277,6 +281,47 @@ def test_stocks_bulk_update_using_variant_external_ref_and_warehouse_id(
     assert stock_2.quantity == new_quantity_2
 
 
+def test_stocks_bulk_update_return_error_when_inputs_have_different_selectors(
+    staff_api_client,
+    variant_with_many_stocks,
+    permission_manage_products,
+):
+    # given
+    variant = variant_with_many_stocks
+    variant_external_reference = variant.external_reference
+
+    stocks = variant.stocks.all()
+    stock_1 = stocks[0]
+    stock_2 = stocks[1]
+
+    warehouse_1_id = graphene.Node.to_global_id("Warehouse", stock_1.warehouse_id)
+
+    stocks_input = [
+        {
+            "variantExternalReference": variant_external_reference,
+            "warehouseId": warehouse_1_id,
+            "quantity": 123,
+        },
+        {
+            "variantExternalReference": variant_external_reference,
+            "warehouseExternalReference": stock_2.warehouse.external_reference,
+            "quantity": 123,
+        },
+    ]
+
+    variables = {"stocks": stocks_input}
+
+    # when
+    staff_api_client.user.user_permissions.add(permission_manage_products)
+    response = staff_api_client.post_graphql(STOCKS_BULK_UPDATE_MUTATION, variables)
+    content = get_graphql_content(response)
+    data = content["data"]["stockBulkUpdate"]
+
+    # then
+    assert data["errors"][0]["code"] == StockBulkUpdateErrorCode.GRAPHQL_ERROR.name
+    assert data["count"] == 0
+
+
 def test_stocks_bulk_update_when_no_variant_args_provided(
     staff_api_client,
     variant_with_many_stocks,
diff --git a/saleor/tests/fixtures.py b/saleor/tests/fixtures.py
--- a/saleor/tests/fixtures.py
+++ b/saleor/tests/fixtures.py
@@ -6750,6 +6750,7 @@ def warehouse_no_shipping_zone(address, channel_USD):
         name="Warehouse without shipping zone",
         slug="warehouse-no-shipping-zone",
         email="test2@example.com",
+        external_reference="warehouse-no-shipping-zone",
     )
     warehouse.channels.add(channel_USD)
     return warehouse