chore(pricing): Fix excessive db queries during price sets update (medusajs#13258)

* chore(pricing): Fix excessive db queries during price sets update

* chore(pricing): Fix excessive db queries during price sets update

* finalize upsert with replace of the rules

* fix limit

* Create quiet-pumpkins-hang.md

---------

Co-authored-by: Oli Juhl <[email protected]>

Author: adrien2p

.changeset/quiet-pumpkins-hang.md

@@ -0,0 +1,5 @@
+---
+"@medusajs/pricing": patch
+---
+
+chore(pricing): Fix excessive db queries during price sets update

integration-tests/http/__tests__/shipping-option/admin/shipping-option.spec.ts

@@ -99,7 +99,7 @@ medusaIntegrationTestRunner({
             `/admin/shipping-option-types`,
             {
               label: "Test",
-              code: 'test',
+              code: "test",
             },
             adminHeaders
           )
@@ -1025,7 +1025,7 @@ medusaIntegrationTestRunner({
           const updateResponse = await api.post(
             `/admin/shipping-options/${shippingOptionId}`,
             {
-              name: "Updated shipping option"
+              name: "Updated shipping option",
             },
             adminHeaders
           )
@@ -1034,7 +1034,7 @@ medusaIntegrationTestRunner({
           expect(updateResponse.data.shipping_option).toEqual(
             expect.objectContaining({
               id: expect.any(String),
-              name: "Updated shipping option"
+              name: "Updated shipping option",
             })
           )
         })
@@ -1129,7 +1129,7 @@ medusaIntegrationTestRunner({
               description: "Test description",
               code: "test-code",
             },
-            type_id: "test_type_id"
+            type_id: "test_type_id",
           }
 
           const error = await api
@@ -1141,7 +1141,9 @@ medusaIntegrationTestRunner({
             .catch((e) => e)
 
           expect(error.response.status).toEqual(400)
-          expect(error.response.data.message).toEqual("Invalid request: Only one of 'type' or 'type_id' can be provided")
+          expect(error.response.data.message).toEqual(
+            "Invalid request: Only one of 'type' or 'type_id' can be provided"
+          )
         })
       })
 

packages/modules/pricing/src/services/pricing-module.ts

@@ -29,6 +29,7 @@ import {
   groupBy,
   InjectManager,
   InjectTransactionManager,
+  isDefined,
   isPresent,
   isString,
   MathBN,
@@ -370,7 +371,6 @@ export default class PricingModuleService
           }
         }
 
-
         pricesSetPricesMap.set(key, { calculatedPrice, originalPrice })
         priceIds.push(
           ...(deduplicate(
@@ -535,6 +535,7 @@ export default class PricingModuleService
     @MedusaContext() sharedContext: Context = {}
   ): Promise<PriceSetDTO | PriceSetDTO[]> {
     const input = Array.isArray(data) ? data : [data]
+
     const forUpdate = input.filter(
       (priceSet): priceSet is ServiceTypes.UpdatePriceSetInput => !!priceSet.id
     )
@@ -621,73 +622,215 @@ export default class PricingModuleService
     data: ServiceTypes.UpdatePriceSetInput[],
     @MedusaContext() sharedContext: Context = {}
   ): Promise<InferEntityType<typeof PriceSet>[]> {
-    // TODO: Since money IDs are rarely passed, this will delete all previous data and insert new entries.
-    // We can make the `insert` inside upsertWithReplace do an `upsert` instead to avoid this
     const normalizedData = await this.normalizeUpdateData(data)
 
-    const priceListPrices = await this.priceService_.list({
-      price_set_id: normalizedData.map(({ id }) => id),
-      price_list_id: { $ne: null },
-    })
+    const priceSetIds = normalizedData.map(({ id }) => id)
+    const existingPrices = await this.priceService_.list(
+      {
+        price_set_id: priceSetIds,
+        price_list_id: null,
+      },
+      {
+        relations: ["price_rules"],
+        take: null,
+      },
+      sharedContext
+    )
+
+    const existingPricesMap = new Map<string, InferEntityType<typeof Price>>(
+      existingPrices.map((p) => [p.id, p])
+    )
 
     const prices = normalizedData.flatMap((priceSet) => priceSet.prices || [])
-    const { entities: upsertedPrices, performedActions } =
-      await this.priceService_.upsertWithReplace(
-        prices,
-        { relations: ["price_rules"] },
+
+    const pricesToCreate = prices.filter(
+      (price) => !price.id || !existingPricesMap.has(price.id)
+    )
+    const pricesToUpdate = prices.filter(
+      (price) => price.id && existingPricesMap.has(price.id)
+    )
+
+    const incomingPriceIds = new Set(prices.map((p) => p.id).filter(Boolean))
+    const pricesToDelete = existingPrices
+      .filter((existingPrice) => !incomingPriceIds.has(existingPrice.id))
+      .map((p) => p.id)
+
+    let createdPrices: InferEntityType<typeof Price>[] = []
+    let updatedPrices: InferEntityType<typeof Price>[] = []
+
+    if (pricesToCreate.length > 0) {
+      createdPrices = await this.priceService_.create(
+        pricesToCreate.map((price) => {
+          price.price_rules ??= []
+          return price
+        }),
         sharedContext
       )
+    }
 
-    composeAllEvents({
-      eventBuilders,
-      performedActions,
-      sharedContext,
-    })
+    if (pricesToUpdate.length > 0) {
+      // Handle price rules for updated prices
+      for (const priceToUpdate of pricesToUpdate) {
+        const existingPrice = existingPricesMap.get(priceToUpdate.id!)
+
+        if (priceToUpdate.price_rules?.length) {
+          const existingPriceRules = existingPrice?.price_rules || []
+
+          // Separate price rules for create, update, delete
+          const priceRulesToCreate = priceToUpdate.price_rules.filter(
+            (rule) => !("id" in rule)
+          )
+          const priceRulesToUpdate = priceToUpdate.price_rules.filter(
+            (rule) => "id" in rule
+          )
+
+          const incomingPriceRuleIds = new Set(
+            priceToUpdate.price_rules
+              .map((r) => "id" in r && r.id)
+              .filter(Boolean)
+          )
+          const priceRulesToDelete = existingPriceRules
+            .filter(
+              (existingRule) => !incomingPriceRuleIds.has(existingRule.id)
+            )
+            .map((r) => r.id)
+
+          let createdPriceRules: InferEntityType<typeof PriceRule>[] = []
+          let updatedPriceRules: InferEntityType<typeof PriceRule>[] = []
+
+          // Bulk operations for price rules
+          if (priceRulesToCreate.length > 0) {
+            createdPriceRules = await this.priceRuleService_.create(
+              priceRulesToCreate.map((rule) => ({
+                ...rule,
+                price_id: priceToUpdate.id,
+              })),
+              sharedContext
+            )
+            eventBuilders.createdPriceRule({
+              data: createdPriceRules.map((r) => ({ id: r.id })),
+              sharedContext,
+            })
+          }
 
-    const priceSetsToUpsert = normalizedData.map((priceSet) => {
-      const { prices, ...rest } = priceSet
-      return {
-        ...rest,
-        prices: [
-          ...upsertedPrices
-            .filter((p) => p.price_set_id === priceSet.id)
-            .map((price) => {
-              // @ts-ignore
-              delete price.price_rules
-              return price
-            }),
-          ...priceListPrices
-            .filter((p) => p.price_set_id === priceSet.id)
-            .map((price) => ({
-              id: price.id,
-              amount: price.amount,
-              price_set_id: price.price_set_id,
-              price_list_id: price.price_list_id,
-            })),
-        ],
+          if (priceRulesToUpdate.length > 0) {
+            updatedPriceRules = await this.priceRuleService_.update(
+              priceRulesToUpdate,
+              sharedContext
+            )
+            eventBuilders.updatedPriceRule({
+              data: updatedPriceRules.map((r) => ({ id: r.id })),
+              sharedContext,
+            })
+          }
+
+          if (priceRulesToDelete.length > 0) {
+            await this.priceRuleService_.delete(
+              priceRulesToDelete,
+              sharedContext
+            )
+            eventBuilders.deletedPriceRule({
+              data: priceRulesToDelete.map((id) => ({ id })),
+              sharedContext,
+            })
+          }
+
+          const upsertedPriceRules = [
+            ...createdPriceRules,
+            ...updatedPriceRules,
+          ]
+
+          priceToUpdate.price_rules = upsertedPriceRules
+          ;(priceToUpdate as InferEntityType<typeof Price>).rules_count =
+            upsertedPriceRules.length
+        } else if (
+          // In the case price_rules is provided but without any rules, we delete the existing rules
+          isDefined(priceToUpdate.price_rules) &&
+          priceToUpdate.price_rules.length === 0
+        ) {
+          const priceRuleToDelete = existingPrice?.price_rules?.map((r) => r.id)
+
+          if (priceRuleToDelete?.length) {
+            await this.priceRuleService_.delete(
+              priceRuleToDelete,
+              sharedContext
+            )
+            eventBuilders.deletedPriceRule({
+              data: priceRuleToDelete.map((id) => ({ id })),
+              sharedContext,
+            })
+          }
+
+          ;(priceToUpdate as InferEntityType<typeof Price>).rules_count = 0
+        } else {
+          // @ts-expect-error - we want to delete the rules_count property in any case even if provided by mistake
+          delete (priceToUpdate as InferEntityType<typeof Price>).rules_count
+        }
+        // We don't want to persist the price_rules in the database through the price service as it would not work
+        delete priceToUpdate.price_rules
       }
-    })
 
-    const { entities: priceSets, performedActions: priceSetPerformedActions } =
-      await this.priceSetService_.upsertWithReplace(
-        priceSetsToUpsert,
-        { relations: ["prices"] },
+      updatedPrices = await this.priceService_.update(
+        pricesToUpdate,
         sharedContext
       )
+    }
 
-    composeAllEvents({
-      eventBuilders,
-      performedActions: priceSetPerformedActions,
-      sharedContext,
+    if (pricesToDelete.length > 0) {
+      await this.priceService_.delete(pricesToDelete, sharedContext)
+    }
+
+    if (createdPrices.length > 0) {
+      eventBuilders.createdPrice({
+        data: createdPrices.map((p) => ({ id: p.id })),
+        sharedContext,
+      })
+    }
+
+    if (updatedPrices.length > 0) {
+      eventBuilders.updatedPrice({
+        data: updatedPrices.map((p) => ({ id: p.id })),
+        sharedContext,
+      })
+    }
+
+    if (pricesToDelete.length > 0) {
+      eventBuilders.deletedPrice({
+        data: pricesToDelete.map((id) => ({ id })),
+        sharedContext,
+      })
+    }
+
+    const priceSets = await this.priceSetService_.list(
+      { id: normalizedData.map(({ id }) => id) },
+      {
+        relations: ["prices", "prices.price_rules"],
+      },
+      sharedContext
+    )
+
+    const upsertedPricesMap = new Map<string, InferEntityType<typeof Price>[]>()
+
+    const upsertedPrices = [...createdPrices, ...updatedPrices]
+    upsertedPrices.forEach((price) => {
+      upsertedPricesMap.set(price.price_set_id, [
+        ...(upsertedPricesMap.get(price.price_set_id) || []),
+        price,
+      ])
     })
 
-    return priceSets.map((ps) => {
-      if (ps.prices) {
-        ps.prices = (ps.prices as any).filter((p) => !p.price_list_id)
-      }
+    // re assign the prices to the price sets to not have to refetch after the transaction and keep the bahaviour the same as expected. If the user needs more data, he can still re list the price set with the expected fields and relations that he needs
+
+    priceSets.forEach((ps) => {
+      ps.prices = upsertedPricesMap.get(ps.id) || []
+    })
 
-      return ps
+    eventBuilders.updatedPriceSet({
+      data: priceSets.map((ps) => ({ id: ps.id })),
+      sharedContext,
     })
+
+    return priceSets
   }
 
   private async normalizeUpdateData(data: ServiceTypes.UpdatePriceSetInput[]) {