From a4e1d69052b2cbef907fae6fc0109809cf5b8d9e Mon Sep 17 00:00:00 2001
From: gmbronco <83549293+gmbronco@users.noreply.github.com>
Date: Thu, 7 Dec 2023 12:18:06 +0100
Subject: [PATCH] V2 options

---
 modules/beethoven/balancer-sdk.gql          |  9 +++++++++
 modules/beethoven/balancer-sdk.resolvers.ts | 22 ++++++++++++++++++++-
 modules/sor/sor.service.ts                  |  5 ++---
 modules/sor/sorV2/sorV2.service.ts          | 18 ++++++++---------
 modules/sor/types.ts                        | 18 ++++++++++++++++-
 5 files changed, 58 insertions(+), 14 deletions(-)

diff --git a/modules/beethoven/balancer-sdk.gql b/modules/beethoven/balancer-sdk.gql
index d26122746..01f5570fe 100644
--- a/modules/beethoven/balancer-sdk.gql
+++ b/modules/beethoven/balancer-sdk.gql
@@ -6,6 +6,7 @@ extend type Query {
         swapType: GqlSorSwapType!
         swapAmount: BigDecimal! #expected in human readable form
         swapOptions: GqlSorSwapOptionsInput!
+        swapOptionsV2: GqlV2SwapOptionsInput
     ): GqlSorGetSwapsResponse!
     sorGetBatchSwapForTokensIn(
         tokensIn: [GqlTokenAmountHumanReadable!]!
@@ -25,6 +26,14 @@ input GqlSorSwapOptionsInput {
     forceRefresh: Boolean #don't use any cached responses
 }
 
+input GqlV2SwapOptionsInput {
+    maxDepth: Int # default 6
+    maxNonBoostedPathDepth: Int # default 3
+    maxNonBoostedHopTokensInBoostedPath: Int # default 2
+    approxPathsToReturn: Int # default 5
+    poolIdsToInclude: [String]
+}
+
 type GqlSorGetSwapsResponse {
     tokenIn: String!
     tokenOut: String!
diff --git a/modules/beethoven/balancer-sdk.resolvers.ts b/modules/beethoven/balancer-sdk.resolvers.ts
index aa830250f..216ffa341 100644
--- a/modules/beethoven/balancer-sdk.resolvers.ts
+++ b/modules/beethoven/balancer-sdk.resolvers.ts
@@ -3,6 +3,9 @@ import { balancerSorService } from './balancer-sor.service';
 import { tokenService } from '../token/token.service';
 import { sorService } from '../sor/sor.service';
 import { getTokenAmountHuman } from '../sor/utils';
+import { SorV2Options } from '../sor/types';
+
+type V2Options = SorV2Options['graphTraversalConfig'];
 
 const balancerSdkResolvers: Resolvers = {
     Query: {
@@ -11,7 +14,24 @@ const balancerSdkResolvers: Resolvers = {
             // Use TokenAmount to help follow scaling requirements in later logic
             // args.swapAmount is HumanScale
             const amount = await getTokenAmountHuman(amountToken, args.swapAmount, args.chain);
-            const swaps = await sorService.getBeetsSwaps({ ...args, swapAmount: amount });
+            const { swapOptionsV2, ...cleanArgs } = args;
+
+            const graphTraversalConfig: V2Options = {};
+            if (swapOptionsV2 && swapOptionsV2.approxPathsToReturn)
+                graphTraversalConfig['approxPathsToReturn'] = swapOptionsV2.approxPathsToReturn;
+            if (swapOptionsV2 && swapOptionsV2.maxDepth) graphTraversalConfig['maxDepth'] = swapOptionsV2.maxDepth;
+            if (swapOptionsV2 && swapOptionsV2.maxNonBoostedHopTokensInBoostedPath)
+                graphTraversalConfig['maxNonBoostedHopTokensInBoostedPath'] =
+                    swapOptionsV2.maxNonBoostedHopTokensInBoostedPath;
+            if (swapOptionsV2 && swapOptionsV2.maxNonBoostedPathDepth)
+                graphTraversalConfig['maxNonBoostedPathDepth'] = swapOptionsV2.maxNonBoostedPathDepth;
+
+            const v2Options = Object.keys(graphTraversalConfig).length > 0 ? { graphTraversalConfig } : undefined;
+            const swaps = await sorService.getBeetsSwaps({
+                ...cleanArgs,
+                swapOptionsV2: v2Options,
+                swapAmount: amount,
+            });
             return { ...swaps, __typename: 'GqlSorGetSwapsResponse' };
         },
         sorGetBatchSwapForTokensIn: async (parent, args, context) => {
diff --git a/modules/sor/sor.service.ts b/modules/sor/sor.service.ts
index 789fb0550..1602e1d34 100644
--- a/modules/sor/sor.service.ts
+++ b/modules/sor/sor.service.ts
@@ -122,9 +122,8 @@ export class SorService {
         v2Time: number,
     ) {
         const sorMetricsPublisher = getSorMetricsPublisher(chain);
-
-        await sorMetricsPublisher.publish(`SOR_VALID_V1`, v1.isValid ? 10 : 1);
-        await sorMetricsPublisher.publish(`SOR_VALID_V2`, v2.isValid ? 10 : 1);
+        await sorMetricsPublisher.publish(`SOR_VALID_V1`, v1.isValid ? 1 : 0);
+        await sorMetricsPublisher.publish(`SOR_VALID_V2`, v2.isValid ? 1 : 0);
 
         if (!version) return;
 
diff --git a/modules/sor/sorV2/sorV2.service.ts b/modules/sor/sorV2/sorV2.service.ts
index 70be4c3ae..851eb20bc 100644
--- a/modules/sor/sorV2/sorV2.service.ts
+++ b/modules/sor/sorV2/sorV2.service.ts
@@ -270,20 +270,20 @@ export class SorV2Service implements SwapService {
         this.cache = new Cache<string, BasePool[]>();
     }
 
-    public async getSwapResult({ chain, tokenIn, tokenOut, swapType, swapAmount }: GetSwapsInput): Promise<SwapResult> {
+    public async getSwapResult({
+        chain,
+        tokenIn,
+        tokenOut,
+        swapType,
+        swapAmount,
+        swapOptionsV2,
+    }: GetSwapsInput): Promise<SwapResult> {
         try {
             const poolsFromDb = await this.getBasePools(chain);
             const tIn = await this.getToken(tokenIn as Address, chain);
             const tOut = await this.getToken(tokenOut as Address, chain);
             const swapKind = this.mapSwapType(swapType);
-            const swap = await sorGetSwapsWithPools(
-                tIn,
-                tOut,
-                swapKind,
-                swapAmount,
-                poolsFromDb,
-                // swapOptions, // I don't think we need specific swapOptions for this?
-            );
+            const swap = await sorGetSwapsWithPools(tIn, tOut, swapKind, swapAmount, poolsFromDb, swapOptionsV2);
             return new SwapResultV2(swap);
         } catch (err: any) {
             console.error(
diff --git a/modules/sor/types.ts b/modules/sor/types.ts
index a46f395c6..8777329dd 100644
--- a/modules/sor/types.ts
+++ b/modules/sor/types.ts
@@ -1,5 +1,11 @@
 import { Chain } from '@prisma/client';
-import { GqlCowSwapApiResponse, GqlSorSwapType, GqlSorGetSwapsResponse, GqlSorSwapOptionsInput } from '../../schema';
+import {
+    GqlCowSwapApiResponse,
+    GqlSorSwapType,
+    GqlSorGetSwapsResponse,
+    GqlSorSwapOptionsInput,
+    GqlV2SwapOptionsInput,
+} from '../../schema';
 import { TokenAmount } from '@balancer/sdk';
 export interface GetSwapsInput {
     chain: Chain;
@@ -8,6 +14,16 @@ export interface GetSwapsInput {
     swapType: GqlSorSwapType;
     swapAmount: TokenAmount;
     swapOptions: GqlSorSwapOptionsInput;
+    swapOptionsV2?: SorV2Options;
+}
+
+export interface SorV2Options {
+    graphTraversalConfig?: {
+        approxPathsToReturn?: number;
+        maxDepth?: number;
+        maxNonBoostedHopTokensInBoostedPath?: number;
+        maxNonBoostedPathDepth?: number;
+    };
 }
 
 export interface SwapResult {