From ad3524ffbe24819a153d947ea0a0266f4760d684 Mon Sep 17 00:00:00 2001
From: Adrien de Peretti <adrien.deperetti@gmail.com>
Date: Mon, 7 Oct 2024 12:07:06 +0200
Subject: [PATCH] feat: Add `useQueryStep` (#9384)

* feat: Query step replacement

* fix

* fix types
---
 packages/core/core-flows/package.json         |  1 +
 .../src/cart/steps/get-variant-price-sets.ts  |  2 +-
 packages/core/core-flows/src/common/index.ts  |  1 +
 .../src/common/steps/use-query-graph.ts       | 32 +++++++++++++++++++
 packages/core/core-flows/tsconfig.spec.json   |  5 ---
 .../src/utils/composer/create-step.ts         |  2 +-
 .../workflows-sdk/src/utils/composer/type.ts  |  2 +-
 yarn.lock                                     |  1 +
 8 files changed, 38 insertions(+), 8 deletions(-)
 create mode 100644 packages/core/core-flows/src/common/steps/use-query-graph.ts
 delete mode 100644 packages/core/core-flows/tsconfig.spec.json

diff --git a/packages/core/core-flows/package.json b/packages/core/core-flows/package.json
index 9fa8154746..84a1af1b56 100644
--- a/packages/core/core-flows/package.json
+++ b/packages/core/core-flows/package.json
@@ -34,6 +34,7 @@
     "@swc/core": "^1.7.28",
     "@swc/jest": "^0.2.36",
     "awilix": "^8.0.1",
+    "expect-type": "^0.20.0",
     "jest": "^29.7.0",
     "pg": "^8.13.0",
     "rimraf": "^5.0.1",
diff --git a/packages/core/core-flows/src/cart/steps/get-variant-price-sets.ts b/packages/core/core-flows/src/cart/steps/get-variant-price-sets.ts
index 5c878fc630..d6cc8d8e38 100644
--- a/packages/core/core-flows/src/cart/steps/get-variant-price-sets.ts
+++ b/packages/core/core-flows/src/cart/steps/get-variant-price-sets.ts
@@ -1,6 +1,6 @@
 import { IPricingModuleService } from "@medusajs/framework/types"
 import { MedusaError, Modules } from "@medusajs/framework/utils"
-import { StepResponse, createStep } from "@medusajs/framework/workflows-sdk"
+import { createStep, StepResponse } from "@medusajs/framework/workflows-sdk"
 
 export interface GetVariantPriceSetsStepInput {
   variantIds: string[]
diff --git a/packages/core/core-flows/src/common/index.ts b/packages/core/core-flows/src/common/index.ts
index caba61e8c0..90f9183dad 100644
--- a/packages/core/core-flows/src/common/index.ts
+++ b/packages/core/core-flows/src/common/index.ts
@@ -3,6 +3,7 @@ export * from "./steps/dismiss-remote-links"
 export * from "./steps/remove-remote-links"
 export * from "./steps/emit-event"
 export * from "./steps/update-remote-links"
+export * from "./steps/use-query-graph"
 export * from "./steps/use-remote-query"
 export * from "./workflows/batch-links"
 export * from "./workflows/create-links"
diff --git a/packages/core/core-flows/src/common/steps/use-query-graph.ts b/packages/core/core-flows/src/common/steps/use-query-graph.ts
new file mode 100644
index 0000000000..2127a3f456
--- /dev/null
+++ b/packages/core/core-flows/src/common/steps/use-query-graph.ts
@@ -0,0 +1,32 @@
+import {
+  GraphResultSet,
+  RemoteJoinerOptions,
+  RemoteQueryFunction,
+  RemoteQueryInput,
+} from "@medusajs/framework/types"
+import { createStep, StepFunction, StepResponse } from "@medusajs/workflows-sdk"
+import { ContainerRegistrationKeys } from "@medusajs/utils"
+
+type EntryStepInput<TEntry extends string> = RemoteQueryInput<TEntry> & {
+  options?: RemoteJoinerOptions
+}
+
+const useQueryGraphStepId = "use-query-graph-step"
+
+const step = createStep(
+  useQueryGraphStepId,
+  async (input: EntryStepInput<any>, { container }) => {
+    const query = container.resolve<RemoteQueryFunction>(
+      ContainerRegistrationKeys.QUERY
+    )
+    const { options, ...queryConfig } = input
+
+    const result = await query.graph(queryConfig as any, options)
+    return new StepResponse(result)
+  }
+)
+
+export const useQueryGraphStep = <const TEntry extends string>(
+  input: EntryStepInput<TEntry>
+): ReturnType<StepFunction<EntryStepInput<TEntry>, GraphResultSet<TEntry>>> =>
+  step(input as any) as unknown as ReturnType<StepFunction<EntryStepInput<TEntry>, GraphResultSet<TEntry>>>
diff --git a/packages/core/core-flows/tsconfig.spec.json b/packages/core/core-flows/tsconfig.spec.json
deleted file mode 100644
index b800dda7ee..0000000000
--- a/packages/core/core-flows/tsconfig.spec.json
+++ /dev/null
@@ -1,5 +0,0 @@
-{
-  "extends": "./tsconfig.json",
-  "include": ["src"],
-  "exclude": ["node_modules", "dist"]
-}
diff --git a/packages/core/workflows-sdk/src/utils/composer/create-step.ts b/packages/core/workflows-sdk/src/utils/composer/create-step.ts
index def2b1ad86..e0607e6fff 100644
--- a/packages/core/workflows-sdk/src/utils/composer/create-step.ts
+++ b/packages/core/workflows-sdk/src/utils/composer/create-step.ts
@@ -6,7 +6,7 @@ import {
 } from "@medusajs/orchestration"
 import { isString, OrchestrationUtils } from "@medusajs/utils"
 import { ulid } from "ulid"
-import { StepResponse, resolveValue } from "./helpers"
+import { resolveValue, StepResponse } from "./helpers"
 import { createStepHandler } from "./helpers/create-step-handler"
 import { proxify } from "./helpers/proxy"
 import {
diff --git a/packages/core/workflows-sdk/src/utils/composer/type.ts b/packages/core/workflows-sdk/src/utils/composer/type.ts
index b10b8b32ec..747d92ea08 100644
--- a/packages/core/workflows-sdk/src/utils/composer/type.ts
+++ b/packages/core/workflows-sdk/src/utils/composer/type.ts
@@ -15,7 +15,7 @@ import { CompensateFn, InvokeFn } from "./create-step"
 export type StepFunctionResult<TOutput extends unknown | unknown[] = unknown> =
   (this: CreateWorkflowComposerContext) => WorkflowData<TOutput>
 
-type StepFunctionReturnConfig<TOutput> = {
+export type StepFunctionReturnConfig<TOutput> = {
   config(
     config: { name?: string } & Omit<
       TransactionStepsDefinition,
diff --git a/yarn.lock b/yarn.lock
index 4f73810a14..fe85860ae5 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -5246,6 +5246,7 @@ __metadata:
     "@swc/core": ^1.7.28
     "@swc/jest": ^0.2.36
     awilix: ^8.0.1
+    expect-type: ^0.20.0
     jest: ^29.7.0
     json-2-csv: ^5.5.4
     pg: ^8.13.0