feat: query.index (#11348)

What: - `query.index` helper. It queries the index module, and aggregate the rest of requested fields/relations if needed like `query.graph`. Not covered in this PR: - Hydrate only sub entities returned by the query. Example: 1 out of 5 variants have returned, it should only hydrate the data of the single entity, currently it will merge all the variants of the product. - Generate types of indexed data example: ```ts const query = container.resolve(ContainerRegistrationKeys.QUERY) await query.index({ entity: "product", fields: [ "id", "description", "status", "variants.sku", "variants.barcode", "variants.material", "variants.options.value", "variants.prices.amount", "variants.prices.currency_code", "variants.inventory_items.inventory.sku", "variants.inventory_items.inventory.description", ], filters: { "variants.sku": { $like: "%-1" }, "variants.prices.amount": { $gt: 30 }, }, pagination: { order: { "variants.prices.amount": "DESC", }, }, }) ``` This query return all products where at least one variant has the title ending in `-1` and at least one price bigger than `30`. The Index Module only hold the data used to paginate and filter, and the returned object is: ```json { "id": "prod_01JKEAM2GJZ14K64R0DHK0JE72", "title": null, "variants": [ { "id": "variant_01JKEAM2HC89GWS95F6GF9C6YA", "sku": "extra-variant-1", "prices": [ { "id": "price_01JKEAM2JADEWWX72F8QDP6QXT", "amount": 80, "currency_code": "USD" } ] } ] } ``` All the rest of the fields will be hydrated from their respective modules, and the final result will be: ```json { "id": "prod_01JKEAY2RJTF8TW9A23KTGY1GD", "description": "extra description", "status": "draft", "variants": [ { "sku": "extra-variant-1", "barcode": null, "material": null, "id": "variant_01JKEAY2S945CRZ6X4QZJ7GVBJ", "options": [ { "value": "Red" } ], "prices": [ { "amount": 20, "currency_code": "CAD", "id": "price_01JKEAY2T2EEYSWZHPGG11B7W7" }, { "amount": 80, "currency_code": "USD", "id": "price_01JKEAY2T2NJK2E5468RK84CAR" } ], "inventory_items": [ { "variant_id": "variant_01JKEAY2S945CRZ6X4QZJ7GVBJ", "inventory_item_id": "iitem_01JKEAY2SNY2AWEHPZN0DDXVW6", "inventory": { "sku": "extra-variant-1", "description": "extra variant 1", "id": "iitem_01JKEAY2SNY2AWEHPZN0DDXVW6" } } ] } ] } ``` Co-authored-by: Adrien de Peretti <25098370+adrien2p@users.noreply.github.com>
2025-02-12 12:55:09 +00:00
parent 8d10731343
commit 22276648ad
19 changed files with 1209 additions and 316 deletions
@@ -133,17 +133,16 @@ export class DataSynchronizer {
        })
      } else {
        // Here we assume that the entity is not indexed anymore as it is not part of the schema object representation and we are cleaning the index
+        // TODO: Drop the partition somewhere
        await promiseAll([
-          this.#indexDataService.delete({
-            selector: {
-              name: entity,
-            },
-          }),
-          this.#indexRelationService.delete({
-            selector: {
-              $or: [{ parent_id: entity }, { child_id: entity }],
-            },
-          }),
+          this.#container.manager.execute(
+            `DELETE FROM "index_data" WHERE "name" = ?`,
+            [entity]
+          ),
+          this.#container.manager.execute(
+            `DELETE FROM "index_relation" WHERE "parent_name" = ? OR "child_name" = ?`,
+            [entity, entity]
+          ),
        ])
      }
    }
@@ -171,14 +170,10 @@ export class DataSynchronizer {
        }
      ),
      this.#updatedStatus(entity, IndexMetadataStatus.PROCESSING),
-      this.#indexDataService.update({
-        data: {
-          staled_at: new Date(),
-        },
-        selector: {
-          name: entity,
-        },
-      }),
+      this.#container.manager.execute(
+        `UPDATE "index_data" SET "staled_at" = NOW() WHERE "name" = ?`,
+        [entity]
+      ),
    ])

    const finalAcknoledgement = await this.syncEntity({
@@ -258,15 +253,10 @@ export class DataSynchronizer {
      entityName
    ] as SchemaObjectEntityRepresentation

-    const { fields, alias, moduleConfig } = schemaEntityObjectRepresentation
+    const { alias, moduleConfig } = schemaEntityObjectRepresentation
    const isLink = !!moduleConfig?.isLink

-    const entityPrimaryKey = fields.find(
-      (field) => !!moduleConfig?.primaryKeys?.includes(field)
-    )
-
-    if (!entityPrimaryKey) {
-      // TODO: for now these are skiped
+    if (!alias) {
      const acknoledgement = {
        lastCursor: pagination.cursor ?? null,
        done: true,
@@ -276,14 +266,27 @@ export class DataSynchronizer {
      return acknoledgement
    }

+    const entityPrimaryKey = "id"
+    const moduleHasId = !!moduleConfig?.primaryKeys?.includes("id")
+    if (!moduleHasId) {
+      const acknoledgement = {
+        lastCursor: pagination.cursor ?? null,
+        err: new Error(
+          "Entity does not have a property 'id'. The 'id' must be provided and must be orderable (e.g ulid)"
+        ),
+      }
+
+      await ack(acknoledgement)
+      return acknoledgement
+    }
+
    let processed = 0
    let currentCursor = pagination.cursor!
    const batchSize = Math.min(pagination.batchSize ?? 100, 100)
    const limit = pagination.limit ?? Infinity
-    let done = false
    let error = null

-    while (processed < limit || !done) {
+    while (processed < limit) {
      const filters: Record<string, any> = {}

      if (currentCursor) {
@@ -306,8 +309,7 @@ export class DataSynchronizer {
        },
      })

-      done = !data.length
-      if (done) {
+      if (!data.length) {
        break
      }

@@ -8,10 +8,12 @@ import {
 import {
  MikroOrmBaseRepository as BaseRepository,
  ContainerRegistrationKeys,
+  deepMerge,
  InjectManager,
  InjectTransactionManager,
  isDefined,
  MedusaContext,
+  promiseAll,
  toMikroORMEntity,
 } from "@medusajs/framework/utils"
 import {
@@ -249,20 +251,29 @@ export class PostgresProvider implements IndexTypes.StorageProvider {

    const select = normalizeFieldsSelection(fields)
    const where = flattenObjectKeys(filters)
+
    const joinWhere = flattenObjectKeys(joinFilters)
    const orderBy = flattenObjectKeys(inputOrderBy)

    const { manager } = sharedContext as { manager: SqlEntityManager }
    let hasPagination = false
-    if (isDefined(skip)) {
+    let hasCount = false
+    if (isDefined(skip) || isDefined(take)) {
      hasPagination = true
+
+      if (isDefined(skip)) {
+        hasCount = true
+      }
    }

+    const requestedFields = deepMerge(deepMerge(select, filters), inputOrderBy)
+
    const connection = manager.getConnection()
    const qb = new QueryBuilder({
      schema: this.schemaObjectRepresentation_,
      entityMap: this.schemaEntitiesMap_,
      knex: connection.getKnex(),
+      rawConfig: config,
      selector: {
        select,
        where,
@@ -274,19 +285,40 @@ export class PostgresProvider implements IndexTypes.StorageProvider {
        keepFilteredEntities,
        orderBy,
      },
+      requestedFields,
    })

-    const sql = qb.buildQuery(hasPagination, !!keepFilteredEntities)
+    const [sql, sqlCount] = qb.buildQuery({
+      hasPagination,
+      returnIdOnly: !!keepFilteredEntities,
+      hasCount,
+    })

-    let resultSet = await manager.execute(sql)
-    const count = hasPagination ? +(resultSet[0]?.count ?? 0) : undefined
+    const promises: Promise<any>[] = []
+
+    promises.push(manager.execute(sql))
+
+    if (hasCount && sqlCount) {
+      promises.push(manager.execute(sqlCount))
+    }
+
+    let [resultSet, count] = await promiseAll(promises)
+
+    const resultMetadata: IndexTypes.QueryFunctionReturnPagination | undefined =
+      hasPagination
+        ? {
+            count: hasCount ? parseInt(count[0].count) : undefined,
+            skip,
+            take,
+          }
+        : undefined

    if (keepFilteredEntities) {
      const mainEntity = Object.keys(select)[0]

      const ids = resultSet.map((r) => r[`${mainEntity}.id`])
      if (ids.length) {
-        return await this.query<TEntry>(
+        const result = await this.query<TEntry>(
          {
            fields,
            joinFilters,
@@ -300,6 +332,8 @@ export class PostgresProvider implements IndexTypes.StorageProvider {
          } as IndexTypes.IndexQueryConfig<TEntry>,
          sharedContext
        )
+        result.metadata ??= resultMetadata
+        return result
      }
    }

@@ -307,13 +341,7 @@ export class PostgresProvider implements IndexTypes.StorageProvider {
      data: qb.buildObjectFromResultset(
        resultSet
      ) as IndexTypes.QueryResultSet<TEntry>["data"],
-      metadata: hasPagination
-        ? {
-            count: count!,
-            skip,
-            take,
-          }
-        : undefined,
+      metadata: resultMetadata,
    }
  }

@@ -365,12 +393,19 @@ export class PostgresProvider implements IndexTypes.StorageProvider {
        return acc
      }, {}) as TData

-      await indexRepository.upsert({
-        id: cleanedEntityData.id,
-        name: entity,
-        data: cleanedEntityData,
-        staled_at: null,
-      })
+      await indexRepository.upsert(
+        {
+          id: cleanedEntityData.id,
+          name: entity,
+          data: cleanedEntityData,
+          staled_at: null,
+        },
+        {
+          onConflictAction: "merge",
+          onConflictFields: ["id", "name"],
+          onConflictMergeFields: ["data", "staled_at"],
+        }
+      )

      /**
       * Retrieve the parents to attach it to the index entry.
@@ -391,12 +426,19 @@ export class PostgresProvider implements IndexTypes.StorageProvider {
          : [parentData]

        for (const parentData_ of parentDataCollection) {
-          await indexRepository.upsert({
-            id: (parentData_ as any).id,
-            name: parentEntity,
-            data: parentData_,
-            staled_at: null,
-          })
+          await indexRepository.upsert(
+            {
+              id: (parentData_ as any).id,
+              name: parentEntity,
+              data: parentData_,
+              staled_at: null,
+            },
+            {
+              onConflictAction: "merge",
+              onConflictFields: ["id", "name"],
+              onConflictMergeFields: ["data", "staled_at"],
+            }
+          )

          await indexRelationRepository.upsert(
            {
@@ -416,6 +458,7 @@ export class PostgresProvider implements IndexTypes.StorageProvider {
                "parent_name",
                "child_name",
              ],
+              onConflictMergeFields: ["staled_at"],
            }
          )
        }
@@ -453,17 +496,24 @@ export class PostgresProvider implements IndexTypes.StorageProvider {
    )

    await indexRepository.upsertMany(
-      data_.map((entityData) => {
-        return {
-          id: entityData.id,
-          name: entity,
-          data: entityProperties.reduce((acc, property) => {
-            acc[property] = entityData[property]
-            return acc
-          }, {}),
-          staled_at: null,
+      data_.map(
+        (entityData) => {
+          return {
+            id: entityData.id,
+            name: entity,
+            data: entityProperties.reduce((acc, property) => {
+              acc[property] = entityData[property]
+              return acc
+            }, {}),
+            staled_at: null,
+          }
+        },
+        {
+          onConflictAction: "merge",
+          onConflictFields: ["id", "name"],
+          onConflictMergeFields: ["data", "staled_at"],
        }
-      })
+      )
    )
  }

@@ -605,12 +655,19 @@ export class PostgresProvider implements IndexTypes.StorageProvider {
        return acc
      }, {}) as TData

-      await indexRepository.upsert({
-        id: cleanedEntityData.id,
-        name: entity,
-        data: cleanedEntityData,
-        staled_at: null,
-      })
+      await indexRepository.upsert(
+        {
+          id: cleanedEntityData.id,
+          name: entity,
+          data: cleanedEntityData,
+          staled_at: null,
+        },
+        {
+          onConflictAction: "merge",
+          onConflictFields: ["id", "name"],
+          onConflictMergeFields: ["data", "staled_at"],
+        }
+      )

      /**
       * Create the index relation entries for the parent entity and the child entity
@@ -634,6 +691,7 @@ export class PostgresProvider implements IndexTypes.StorageProvider {
            "parent_name",
            "child_name",
          ],
+          onConflictMergeFields: ["staled_at"],
        }
      )

@@ -655,6 +713,7 @@ export class PostgresProvider implements IndexTypes.StorageProvider {
            "parent_name",
            "child_name",
          ],
+          onConflictMergeFields: ["staled_at"],
        }
      )
    }
@@ -1,6 +1,6 @@
+import { IndexTypes } from "@medusajs/framework/types"
 import { SqlEntityManager } from "@mikro-orm/postgresql"
 import { schemaObjectRepresentationPropertiesToOmit } from "@types"
-import { IndexTypes } from "@medusajs/framework/types"

 export async function createPartitions(
  schemaObjectRepresentation: IndexTypes.SchemaObjectRepresentation,
@@ -54,6 +54,10 @@ export async function createPartitions(
        `CREATE INDEX CONCURRENTLY IF NOT EXISTS "IDX_cat_${cName}_data_gin" ON ${activeSchema}cat_${cName} USING GIN ("data" jsonb_path_ops)`
      )

+      part.push(
+        `CREATE INDEX CONCURRENTLY IF NOT EXISTS "IDX_cat_${cName}_id" ON ${activeSchema}cat_${cName} ("id")`
+      )
+
      // create child id index on pivot partitions
      for (const parent of schemaObjectRepresentation[key].parents) {
        const pName = `${parent.ref.entity}${key}`.toLowerCase()
@@ -1,5 +1,10 @@
 import { IndexTypes } from "@medusajs/framework/types"
-import { GraphQLUtils, isObject, isString } from "@medusajs/framework/utils"
+import {
+  GraphQLUtils,
+  isObject,
+  isPresent,
+  isString,
+} from "@medusajs/framework/utils"
 import { Knex } from "@mikro-orm/knex"
 import { OrderBy, QueryFormat, QueryOptions, Select } from "@types"

@@ -24,6 +29,10 @@ export class QueryBuilder {
  private readonly options?: QueryOptions
  private readonly schema: IndexTypes.SchemaObjectRepresentation
  private readonly allSchemaFields: Set<string>
+  private readonly rawConfig?: IndexTypes.IndexQueryConfig<any>
+  private readonly requestedFields: {
+    [key: string]: any
+  }

  constructor(args: {
    schema: IndexTypes.SchemaObjectRepresentation
@@ -31,6 +40,10 @@ export class QueryBuilder {
    knex: Knex
    selector: QueryFormat
    options?: QueryOptions
+    rawConfig?: IndexTypes.IndexQueryConfig<any>
+    requestedFields: {
+      [key: string]: any
+    }
  }) {
    this.schema = args.schema
    this.entityMap = args.entityMap
@@ -41,6 +54,8 @@ export class QueryBuilder {
    this.allSchemaFields = new Set(
      Object.values(this.schema).flatMap((entity) => entity.fields ?? [])
    )
+    this.rawConfig = args.rawConfig
+    this.requestedFields = args.requestedFields
  }

  private getStructureKeys(structure) {
@@ -56,7 +71,9 @@ export class QueryBuilder {
        return
      }

-      throw new Error(`Could not find entity for path: ${path}`)
+      throw new Error(
+        `Could not find entity for path: ${path}. It might not be indexed.`
+      )
    }

    return this.schema._schemaPropertiesMap[path]
@@ -66,7 +83,7 @@ export class QueryBuilder {
    const entity = this.getEntity(path)?.ref?.entity!
    const fieldRef = this.entityMap[entity]._fields[field]
    if (!fieldRef) {
-      throw new Error(`Field ${field} not found in the entityMap.`)
+      throw new Error(`Field ${field} is not indexed.`)
    }

    let currentType = fieldRef.type
@@ -224,6 +241,8 @@ export class QueryBuilder {
            const val = operator === "IN" ? subValue : [subValue]
            if (operator === "=" && subValue === null) {
              operator = "IS"
+            } else if (operator === "!=" && subValue === null) {
+              operator = "IS NOT"
            }

            if (operator === "=") {
@@ -306,13 +325,16 @@ export class QueryBuilder {

    const isSelectableField = this.allSchemaFields.has(parentProperty)
    const entities = this.getEntity(currentAliasPath, false)
-    if (isSelectableField || !entities) {
+    const entityRef = entities?.ref!
+
+    // !entityRef.alias means the object has not table, it's a nested object
+    if (isSelectableField || !entities || !entityRef?.alias) {
      // We are currently selecting a specific field of the parent entity or the entity is not found on the index schema
      // We don't need to build the query parts for this as there is no join
      return []
    }

-    const mainEntity = entities.ref.entity
+    const mainEntity = entityRef.entity
    const mainAlias =
      this.getShortAlias(aliasMapping, mainEntity.toLowerCase()) + level

@@ -530,10 +552,18 @@ export class QueryBuilder {
    return result
  }

-  public buildQuery(countAllResults = true, returnIdOnly = false): string {
+  public buildQuery({
+    hasPagination = true,
+    hasCount = false,
+    returnIdOnly = false,
+  }: {
+    hasPagination?: boolean
+    hasCount?: boolean
+    returnIdOnly?: boolean
+  }): [string, string | null] {
    const queryBuilder = this.knex.queryBuilder()

-    const structure = this.structure
+    const structure = this.requestedFields
    const filter = this.selector.where ?? {}

    const { orderBy: order, skip, take } = this.options ?? {}
@@ -564,15 +594,6 @@ export class QueryBuilder {
      ? this.buildSelectParts(rootStructure, rootKey, aliasMapping)
      : { [rootKey + ".id"]: `${rootAlias}.id` }

-    if (countAllResults) {
-      selectParts["offset_"] = this.knex.raw(
-        `DENSE_RANK() OVER (ORDER BY ${this.getShortAlias(
-          aliasMapping,
-          rootEntity
-        )}.id)`
-      )
-    }
-
    queryBuilder.select(selectParts)

    queryBuilder.from(
@@ -601,24 +622,150 @@ export class QueryBuilder {
      )
    }

-    let sql = `WITH data AS (${queryBuilder.toQuery()})
-    SELECT * ${
-      countAllResults ? ", (SELECT max(offset_) FROM data) AS count" : ""
-    }
-    FROM data`
+    let distinctQueryBuilder = queryBuilder.clone()

    let take_ = !isNaN(+take!) ? +take! : 15
    let skip_ = !isNaN(+skip!) ? +skip! : 0
-    if (typeof take === "number" || typeof skip === "number") {
-      sql += `
-        WHERE offset_ > ${skip_}
-          AND offset_ <= ${skip_ + take_}
-      `
+    let sql = ""
+
+    if (hasPagination) {
+      const idColumn = `${this.getShortAlias(aliasMapping, rootEntity)}.id`
+      distinctQueryBuilder.clearSelect()
+      distinctQueryBuilder.select(
+        this.knex.raw(`DISTINCT ON (${idColumn}) ${idColumn} as "id"`)
+      )
+      distinctQueryBuilder.limit(take_)
+      distinctQueryBuilder.offset(skip_)
+
+      sql += `WITH paginated_data AS (${distinctQueryBuilder.toQuery()}),`
+
+      queryBuilder.andWhere(
+        this.knex.raw(`${idColumn} IN (SELECT id FROM "paginated_data")`)
+      )
    }

-    return sql
+    sql += `${hasPagination ? " " : "WITH"} data AS (${queryBuilder.toQuery()})
+    SELECT * 
+    FROM data`
+
+    let sqlCount = ""
+    if (hasCount) {
+      sqlCount = this.buildQueryCount()
+    }
+
+    return [sql, hasCount ? sqlCount : null]
  }

+  public buildQueryCount(): string {
+    const queryBuilder = this.knex.queryBuilder()
+
+    const hasWhere = isPresent(this.rawConfig?.filters)
+    const structure = hasWhere ? this.rawConfig?.filters! : this.requestedFields
+
+    const rootKey = this.getStructureKeys(structure)[0]
+
+    const rootStructure = structure[rootKey] as Select
+
+    const entity = this.getEntity(rootKey)!.ref.entity
+    const rootEntity = entity.toLowerCase()
+    const aliasMapping: { [path: string]: string } = {}
+
+    const joinParts = this.buildQueryParts(
+      rootStructure,
+      "",
+      entity,
+      rootKey,
+      [],
+      0,
+      aliasMapping
+    )
+
+    const rootAlias = aliasMapping[rootKey]
+
+    queryBuilder.select(
+      this.knex.raw(`COUNT(DISTINCT ${rootAlias}.id) as count`)
+    )
+
+    queryBuilder.from(
+      `cat_${rootEntity} AS ${this.getShortAlias(aliasMapping, rootEntity)}`
+    )
+
+    if (hasWhere) {
+      joinParts.forEach((joinPart) => {
+        queryBuilder.joinRaw(joinPart)
+      })
+
+      this.parseWhere(aliasMapping, this.selector.where!, queryBuilder)
+    }
+
+    return queryBuilder.toQuery()
+  }
+
+  // NOTE: We are keeping the bellow code for now as reference to alternative implementation for us. DO NOT REMOVE
+  // public buildQueryCount(): string {
+  //   const queryBuilder = this.knex.queryBuilder()
+
+  //   const hasWhere = isPresent(this.rawConfig?.filters)
+  //   const structure = hasWhere ? this.rawConfig?.filters! : this.structure
+
+  //   const rootKey = this.getStructureKeys(structure)[0]
+
+  //   const rootStructure = structure[rootKey] as Select
+
+  //   const entity = this.getEntity(rootKey)!.ref.entity
+  //   const rootEntity = entity.toLowerCase()
+  //   const aliasMapping: { [path: string]: string } = {}
+
+  //   const joinParts = this.buildQueryParts(
+  //     rootStructure,
+  //     "",
+  //     entity,
+  //     rootKey,
+  //     [],
+  //     0,
+  //     aliasMapping
+  //   )
+
+  //   const rootAlias = aliasMapping[rootKey]
+
+  //   queryBuilder.select(this.knex.raw(`COUNT(${rootAlias}.id) as count`))
+
+  //   queryBuilder.from(
+  //     `cat_${rootEntity} AS ${this.getShortAlias(aliasMapping, rootEntity)}`
+  //   )
+
+  //   const self = this
+  //   if (hasWhere && joinParts.length) {
+  //     const fromExistsRaw = joinParts.shift()!
+  //     const [joinPartsExists, fromExistsPart] =
+  //       fromExistsRaw.split(" left join ")
+  //     const [fromExists, whereExists] = fromExistsPart.split(" on ")
+  //     joinParts.unshift(joinPartsExists)
+
+  //     queryBuilder.whereExists(function () {
+  //       this.select(self.knex.raw(`1`))
+  //       this.from(self.knex.raw(`${fromExists}`))
+  //       this.joinRaw(joinParts.join("\n"))
+  //       if (hasWhere) {
+  //         self.parseWhere(aliasMapping, self.selector.where!, this)
+  //         this.whereRaw(self.knex.raw(whereExists))
+  //         return
+  //       }
+
+  //       this.whereRaw(self.knex.raw(whereExists))
+  //     })
+  //   } else {
+  //     queryBuilder.whereExists(function () {
+  //       this.select(self.knex.raw(`1`))
+  //       if (hasWhere) {
+  //         self.parseWhere(aliasMapping, self.selector.where!, this)
+  //       }
+  //     })
+  //   }
+
+  //   return queryBuilder.toQuery()
+  // }
+
  public buildObjectFromResultset(
    resultSet: Record<string, any>[]
  ): Record<string, any>[] {