feat(product): build variant images when retrieving product (#13731)

**What**
-  build variant images when retrieving product

Co-authored-by: Oli Juhl <59018053+olivermrbl@users.noreply.github.com>

.changeset/tame-fireants-compare.md

@@ -0,0 +1,5 @@
+---
+"@medusajs/product": patch
+---
+
+feat(product): build variant images when retrieving product

packages/modules/product/integration-tests/__tests__/product-module-service/products.spec.ts

@@ -3,11 +3,7 @@ import {
   ProductCategoryDTO,
   ProductTagDTO,
 } from "@medusajs/framework/types"
-import {
+import { kebabCase, Modules, ProductStatus } from "@medusajs/framework/utils"
-  kebabCase,
-  Modules,
-  ProductStatus,
-} from "@medusajs/framework/utils"
 import {
   Product,
   ProductCategory,
@@ -16,9 +12,7 @@ import {
   ProductType,
 } from "@models"
 
-import {
+import { moduleIntegrationTestRunner } from "@medusajs/test-utils"
-  moduleIntegrationTestRunner,
-} from "@medusajs/test-utils"
 import { UpdateProductInput } from "@types"
 import {
   buildProductAndRelationsData,
@@ -566,7 +560,6 @@ moduleIntegrationTestRunner<IProductModuleService>({
           )
         })
 
-
         it("should add relationships to a product", async () => {
           const updateData = {
             id: productOne.id,
@@ -1053,7 +1046,6 @@ moduleIntegrationTestRunner<IProductModuleService>({
           )
         })
 
-
         it("should throw because variant doesn't have all options set", async () => {
           const error = await service
             .createProducts([
@@ -1232,7 +1224,6 @@ moduleIntegrationTestRunner<IProductModuleService>({
 
           expect(softDeleted).toHaveLength(1)
         })
-
       })
 
       describe("restore", function () {
@@ -1569,6 +1560,114 @@ moduleIntegrationTestRunner<IProductModuleService>({
             }),
           ])
         })
+
+        it("should populate variant.images when variants.images relation is requested", async () => {
+          const images = [
+            { url: "general-image-1" },
+            { url: "general-image-2" },
+            { url: "variant-specific-image" },
+          ]
+
+          const [product] = await service.createProducts([
+            buildProductAndRelationsData({
+              images,
+              options: [{ title: "size", values: ["small", "large"] }],
+              variants: [
+                { title: "Small", options: { size: "small" } },
+                { title: "Large", options: { size: "large" } },
+              ],
+            }),
+          ])
+
+          const generalImage1 = product.images.find(
+            (img) => img.url === "general-image-1"
+          )!
+          const generalImage2 = product.images.find(
+            (img) => img.url === "general-image-2"
+          )!
+          const variantSpecificImage = product.images.find(
+            (img) => img.url === "variant-specific-image"
+          )!
+
+          const smallVariant = product.variants.find(
+            (v) => v.title === "Small"
+          )!
+          const largeVariant = product.variants.find(
+            (v) => v.title === "Large"
+          )!
+
+          // Add variant-specific image assignment
+          await service.addImageToVariant([
+            {
+              image_id: variantSpecificImage.id,
+              variant_id: smallVariant.id,
+            },
+          ])
+
+          // Test retrieveProduct with variants.images relation
+          const retrievedProduct = await service.retrieveProduct(product.id, {
+            relations: ["variants", "variants.images", "images"],
+          })
+
+          expect(retrievedProduct.variants).toHaveLength(2)
+
+          // First variant (Small) should have general images + variant-specific image
+          const retrievedSmallVariant = retrievedProduct.variants.find(
+            (v) => v.title === "Small"
+          )!
+          expect(retrievedSmallVariant.images).toHaveLength(3) // 2 general + 1 variant-specific
+          expect(retrievedSmallVariant.images).toEqual(
+            expect.arrayContaining([
+              expect.objectContaining({ id: generalImage1.id }),
+              expect.objectContaining({ id: generalImage2.id }),
+              expect.objectContaining({ id: variantSpecificImage.id }),
+            ])
+          )
+
+          // Second variant (Large) should have only general images
+          const retrievedLargeVariant = retrievedProduct.variants.find(
+            (v) => v.title === "Large"
+          )!
+          expect(retrievedLargeVariant.images).toHaveLength(2) // 2 general images only
+          expect(retrievedLargeVariant.images).toEqual(
+            expect.arrayContaining([
+              expect.objectContaining({ id: generalImage1.id }),
+              expect.objectContaining({ id: generalImage2.id }),
+            ])
+          )
+
+          // Test listProducts with variants.images relation
+          const products = await service.listProducts(
+            { id: product.id },
+            { relations: ["variants", "variants.images", "images"] }
+          )
+
+          expect(products).toHaveLength(1)
+          expect(products[0].variants).toHaveLength(2)
+
+          const listSmallVariant = products[0].variants.find(
+            (v) => v.title === "Small"
+          )!
+          expect(listSmallVariant.images).toHaveLength(3)
+          expect(listSmallVariant.images).toEqual(
+            expect.arrayContaining([
+              expect.objectContaining({ id: generalImage1.id }),
+              expect.objectContaining({ id: generalImage2.id }),
+              expect.objectContaining({ id: variantSpecificImage.id }),
+            ])
+          )
+
+          const listLargeVariant = products[0].variants.find(
+            (v) => v.title === "Large"
+          )!
+          expect(listLargeVariant.images).toHaveLength(2)
+          expect(listLargeVariant.images).toEqual(
+            expect.arrayContaining([
+              expect.objectContaining({ id: generalImage1.id }),
+              expect.objectContaining({ id: generalImage2.id }),
+            ])
+          )
+        })
       })
     })
   },

packages/modules/product/src/services/product-module-service.ts

@@ -41,6 +41,7 @@ import {
   MedusaService,
   MessageAggregator,
   Modules,
+  partitionArray,
   ProductStatus,
   removeUndefined,
   toHandle,
@@ -199,12 +200,30 @@ export default class ProductModuleService
     config?: FindConfig<ProductTypes.ProductDTO>,
     @MedusaContext() sharedContext?: Context
   ): Promise<ProductTypes.ProductDTO> {
+    const relationsSet = new Set(config?.relations ?? [])
+    const shouldLoadVariantImages = relationsSet.has("variants.images")
+    if (shouldLoadVariantImages) {
+      relationsSet.add("variants")
+      relationsSet.add("images")
+    }
+
     const product = await this.productService_.retrieve(
       productId,
-      this.getProductFindConfig_(config),
+      this.getProductFindConfig_({
+        ...config,
+        relations: Array.from(relationsSet),
+      }),
       sharedContext
     )
 
+    if (shouldLoadVariantImages && product.variants && product.images) {
+      await this.buildVariantImagesFromProduct(
+        product.variants,
+        product.images,
+        sharedContext
+      )
+    }
+
     return this.baseRepository_.serialize<ProductTypes.ProductDTO>(product)
   }
 
@@ -215,12 +234,34 @@ export default class ProductModuleService
     config?: FindConfig<ProductTypes.ProductDTO>,
     sharedContext?: Context
   ): Promise<ProductTypes.ProductDTO[]> {
+    const relationsSet = new Set(config?.relations ?? [])
+    const shouldLoadVariantImages = relationsSet.has("variants.images")
+    if (shouldLoadVariantImages) {
+      relationsSet.add("variants")
+      relationsSet.add("images")
+    }
+
     const products = await this.productService_.list(
       filters,
-      this.getProductFindConfig_(config),
+      this.getProductFindConfig_({
+        ...config,
+        relations: Array.from(relationsSet),
+      }),
       sharedContext
     )
 
+    if (shouldLoadVariantImages) {
+      for (const product of products) {
+        if (product.variants && product.images) {
+          await this.buildVariantImagesFromProduct(
+            product.variants,
+            product.images,
+            sharedContext
+          )
+        }
+      }
+    }
+
     return this.baseRepository_.serialize<ProductTypes.ProductDTO[]>(products)
   }
 
@@ -231,12 +272,38 @@ export default class ProductModuleService
     config?: FindConfig<ProductTypes.ProductDTO>,
     sharedContext?: Context
   ): Promise<[ProductTypes.ProductDTO[], number]> {
+    const shouldLoadVariantImages =
+      config?.relations?.includes("variants.images")
+
+    // Ensure we load necessary relations
+    const relations = [...(config?.relations || [])]
+    if (shouldLoadVariantImages) {
+      if (!relations.includes("variants")) {
+        relations.push("variants")
+      }
+      if (!relations.includes("images")) {
+        relations.push("images")
+      }
+    }
+
     const [products, count] = await this.productService_.listAndCount(
       filters,
-      this.getProductFindConfig_(config),
+      this.getProductFindConfig_({ ...config, relations }),
       sharedContext
     )
 
+    if (shouldLoadVariantImages) {
+      for (const product of products) {
+        if (product.variants && product.images) {
+          await this.buildVariantImagesFromProduct(
+            product.variants,
+            product.images,
+            sharedContext
+          )
+        }
+      }
+    }
+
     const serializedProducts = await this.baseRepository_.serialize<
       ProductTypes.ProductDTO[]
     >(products)
@@ -2450,4 +2517,57 @@ export default class ProductModuleService
 
     return result
   }
+
+  private async buildVariantImagesFromProduct(
+    variants: InferEntityType<typeof ProductVariant>[],
+    productImages: InferEntityType<typeof ProductImage>[],
+    sharedContext: Context = {}
+  ): Promise<void> {
+    // Create a clean map of images without problematic collections
+    const imagesMap = new Map<string, InferEntityType<typeof ProductImage>>()
+    for (const img of productImages) {
+      imagesMap.set(img.id, img)
+    }
+
+    const variantIds = variants.map((v) => v.id)
+    const variantImageRelations =
+      await this.productVariantProductImageService_.list(
+        { variant_id: variantIds },
+        { select: ["variant_id", "image_id"] },
+        sharedContext
+      )
+
+    const variantIdImageIdsMap = new Map<string, string[]>()
+    const imageIdVariantIdsMap = new Map<string, string[]>()
+
+    // build both lookup
+    for (const relation of variantImageRelations) {
+      if (!variantIdImageIdsMap.has(relation.variant_id)) {
+        variantIdImageIdsMap.set(relation.variant_id, [])
+      }
+      variantIdImageIdsMap.get(relation.variant_id)!.push(relation.image_id)
+
+      if (!imageIdVariantIdsMap.has(relation.image_id)) {
+        imageIdVariantIdsMap.set(relation.image_id, [])
+      }
+      imageIdVariantIdsMap.get(relation.image_id)!.push(relation.variant_id)
+    }
+
+    const [generalImages, variantImages] = partitionArray(
+      productImages,
+      (img) => !imageIdVariantIdsMap.has(img.id) // if image doesn't have any variants, it is a general image
+    )
+
+    for (const variant of variants) {
+      const variantImageIds = variantIdImageIdsMap.get(variant.id) || []
+
+      variant.images = [...generalImages]
+
+      for (const image of variantImages) {
+        if (variantImageIds.includes(image.id)) {
+          variant.images.push(image)
+        }
+      }
+    }
+  }
 }