From 98387231927e9872f54c9e72597576f3273de506 Mon Sep 17 00:00:00 2001
From: Philip Korsholm <88927411+pKorsholm@users.noreply.github.com>
Date: Wed, 3 May 2023 11:17:55 +0200
Subject: [PATCH] fix(medusa): Performance improvement for variant availability
 (#3921)

**What**
- refactor availability invocations to only do one call to the inventory module

@adrien2p would love to run through some tests of this with you and get your take on approach
---
 .changeset/lazy-seas-wait.md                  |   5 +
 .../inventory/products/list-variants.js       |   1 +
 .../src/services/product-variant-inventory.ts | 114 ++++++++++++++----
 3 files changed, 98 insertions(+), 22 deletions(-)
 create mode 100644 .changeset/lazy-seas-wait.md

diff --git a/.changeset/lazy-seas-wait.md b/.changeset/lazy-seas-wait.md
new file mode 100644
index 0000000000..0ea437baa4
--- /dev/null
+++ b/.changeset/lazy-seas-wait.md
@@ -0,0 +1,5 @@
+---
+"@medusajs/medusa": patch
+---
+
+fix(medusa): update query performance for getting product and variant avaiability
diff --git a/integration-tests/plugins/__tests__/inventory/products/list-variants.js b/integration-tests/plugins/__tests__/inventory/products/list-variants.js
index 42c617f175..2fc403838c 100644
--- a/integration-tests/plugins/__tests__/inventory/products/list-variants.js
+++ b/integration-tests/plugins/__tests__/inventory/products/list-variants.js
@@ -91,6 +91,7 @@ describe("List Variants", () => {
         stocked_quantity: 10,
       })
     })
+
     it("Decorates inventory quantities when listing variants", async () => {
       const api = useApi()
 
diff --git a/packages/medusa/src/services/product-variant-inventory.ts b/packages/medusa/src/services/product-variant-inventory.ts
index d6a1f82156..e545086fcf 100644
--- a/packages/medusa/src/services/product-variant-inventory.ts
+++ b/packages/medusa/src/services/product-variant-inventory.ts
@@ -5,6 +5,7 @@ import {
   IInventoryService,
   IStockLocationService,
   InventoryItemDTO,
+  InventoryLevelDTO,
   ReservationItemDTO,
   ReserveQuantityContext,
 } from "@medusajs/types"
@@ -27,6 +28,11 @@ type InjectedDependencies = {
   eventBusService: IEventBusService
 }
 
+type AvailabilityContext = {
+  variantInventoryMap?: Map<string, ProductVariantInventoryItem[]>
+  inventoryLocationMap?: Map<string, InventoryLevelDTO[]>
+}
+
 class ProductVariantInventoryService extends TransactionBaseService {
   protected manager_: EntityManager
   protected transactionManager_: EntityManager | undefined
@@ -652,25 +658,19 @@ class ProductVariantInventoryService extends TransactionBaseService {
   async setVariantAvailability(
     variants: ProductVariant[] | PricedVariant[],
     salesChannelId: string | string[] | undefined,
-    variantInventoryMap: Map<string, ProductVariantInventoryItem[]> = new Map()
+    availabilityContext: AvailabilityContext = {}
   ): Promise<ProductVariant[] | PricedVariant[]> {
     if (!this.inventoryService_) {
       return variants
     }
 
-    if (!variantInventoryMap.size) {
-      const variantInventories = await this.listByVariant(
-        variants.map((v) => v.id)
+    const { variantInventoryMap, inventoryLocationMap } =
+      await this.getAvailabilityContext(
+        variants.map((v) => v.id),
+        salesChannelId,
+        availabilityContext
       )
 
-      variantInventories.forEach((inventory) => {
-        const variantId = inventory.variant_id
-        const currentInventories = variantInventoryMap.get(variantId) || []
-        currentInventories.push(inventory)
-        variantInventoryMap.set(variantId, currentInventories)
-      })
-    }
-
     return await Promise.all(
       variants.map(async (variant) => {
         if (!variant.id) {
@@ -698,15 +698,8 @@ class ProductVariantInventoryService extends TransactionBaseService {
           return variant
         }
 
-        const locationIds =
-          await this.salesChannelLocationService_.listLocationIds(
-            salesChannelId
-          )
-
-        const [locations] = await this.inventoryService_.listInventoryLevels({
-          location_id: locationIds,
-          inventory_item_id: variantInventory[0].inventory_item_id,
-        })
+        const locations =
+          inventoryLocationMap.get(variantInventory[0].inventory_item_id) ?? []
 
         variant.inventory_quantity = locations.reduce(
           (acc, next) => acc + (next.stocked_quantity - next.reserved_quantity),
@@ -721,10 +714,86 @@ class ProductVariantInventoryService extends TransactionBaseService {
     )
   }
 
+  private async getAvailabilityContext(
+    variants: string[],
+    salesChannelId: string | string[] | undefined,
+    existingContext: AvailabilityContext = {}
+  ): Promise<Required<AvailabilityContext>> {
+    let variantInventoryMap = existingContext.variantInventoryMap
+    let inventoryLocationMap = existingContext.inventoryLocationMap
+
+    if (!variantInventoryMap) {
+      variantInventoryMap = new Map()
+      const variantInventories = await this.listByVariant(variants)
+
+      variantInventories.forEach((inventory) => {
+        const variantId = inventory.variant_id
+        const currentInventories = variantInventoryMap!.get(variantId) || []
+        currentInventories.push(inventory)
+        variantInventoryMap!.set(variantId, currentInventories)
+      })
+    }
+
+    const locationIds: string[] = []
+
+    if (salesChannelId && !inventoryLocationMap) {
+      const locations = await this.salesChannelLocationService_
+        .withTransaction(this.activeManager_)
+        .listLocationIds(salesChannelId)
+      locationIds.push(...locations)
+    }
+
+    if (!inventoryLocationMap) {
+      inventoryLocationMap = new Map()
+    }
+
+    if (locationIds.length) {
+      const [locationLevels] = await this.inventoryService_.listInventoryLevels(
+        {
+          location_id: locationIds,
+          inventory_item_id: [
+            ...new Set(
+              Array.from(variantInventoryMap.values())
+                .flat()
+                .map((i) => i.inventory_item_id)
+            ),
+          ],
+        },
+        {},
+        {
+          transactionManager: this.activeManager_,
+        }
+      )
+
+      locationLevels.reduce((acc, curr) => {
+        if (!acc.has(curr.inventory_item_id)) {
+          acc.set(curr.inventory_item_id, [])
+        }
+        acc.get(curr.inventory_item_id)!.push(curr)
+
+        return acc
+      }, inventoryLocationMap)
+    }
+
+    return {
+      variantInventoryMap,
+      inventoryLocationMap,
+    }
+  }
+
   async setProductAvailability(
     products: (Product | PricedProduct)[],
     salesChannelId: string | string[] | undefined
   ): Promise<(Product | PricedProduct)[]> {
+    const variantIds: string[] = products
+      .flatMap((p) => p.variants.map((v: { id?: string }) => v.id) ?? [])
+      .filter((v): v is string => !!v)
+
+    const availabilityContext = await this.getAvailabilityContext(
+      variantIds,
+      salesChannelId
+    )
+
     return await Promise.all(
       products.map(async (product) => {
         if (!product.variants || product.variants.length === 0) {
@@ -733,7 +802,8 @@ class ProductVariantInventoryService extends TransactionBaseService {
 
         product.variants = await this.setVariantAvailability(
           product.variants,
-          salesChannelId
+          salesChannelId,
+          availabilityContext
         )
 
         return product