feat(SQS): custom quote query (#7001) (#7024)

* feat(SQS): custom quote query

* lint

* add test for ParseNumbers

(cherry picked from commit 8214a60)

Co-authored-by: Roman <roman@osmosis.team>

ingest/sqs/domain/mvc/router.go

@@ -36,6 +36,10 @@ type RouterUsecase interface {
 	GetOptimalQuote(ctx context.Context, tokenIn sdk.Coin, tokenOutDenom string) (domain.Quote, error)
 	// GetBestSingleRouteQuote returns the best single route quote for the given tokenIn and tokenOutDenom.
 	GetBestSingleRouteQuote(ctx context.Context, tokenIn sdk.Coin, tokenOutDenom string) (domain.Quote, error)
+	// GetCustomQuote returns the custom quote for the given tokenIn, tokenOutDenom and poolIDs.
+	// It searches for the route that contains the specified poolIDs in the given order.
+	// If such route is not found it returns an error.
+	GetCustomQuote(ctx context.Context, tokenIn sdk.Coin, tokenOutDenom string, poolIDs []uint64) (domain.Quote, error)
 	// GetCandidateRoutes returns the candidate routes for the given tokenIn and tokenOutDenom.
 	GetCandidateRoutes(ctx context.Context, tokenInDenom, tokenOutDenom string) (route.CandidateRoutes, error)
 	// StoreRoutes stores all router state in the files locally. Used for debugging.

ingest/sqs/router/delivery/http/export_test.go

@@ -0,0 +1,5 @@
+package http
+
+func ParseNumbers(numbersParam string) ([]uint64, error) {
+	return parseNumbers(numbersParam)
+}

ingest/sqs/router/delivery/http/router_handler.go

@@ -4,6 +4,8 @@ import (
 	"errors"
 	"net/http"
 	"regexp"
+	"strconv"
+	"strings"
 
 	"github.com/labstack/echo"
 	"github.com/sirupsen/logrus"
@@ -41,6 +43,7 @@ func NewRouterHandler(e *echo.Echo, us mvc.RouterUsecase, logger log.Logger) {
 	e.GET("/quote", handler.GetOptimalQuote)
 	e.GET("/single-quote", handler.GetBestSingleRouteQuote)
 	e.GET("/routes", handler.GetCandidateRoutes)
+	e.GET("/custom-quote", handler.GetCustomQuote)
 	e.POST("/store-state", handler.StoreRouterStateInFiles)
 }
 
@@ -88,6 +91,37 @@ func (a *RouterHandler) GetBestSingleRouteQuote(c echo.Context) error {
 	return c.JSON(http.StatusOK, quote)
 }
 
+// GetCustomQuote returns a direct custom quote. It ensures that the route contains all the pools
+// listed in the specific order, returns error if such route is not found.
+func (a *RouterHandler) GetCustomQuote(c echo.Context) error {
+	ctx := c.Request().Context()
+
+	tokenOutDenom, tokenIn, err := getValidRoutingParameters(c)
+	if err != nil {
+		return err
+	}
+
+	poolIDsStr := c.QueryParam("poolIDs")
+	if len(poolIDsStr) == 0 {
+		return c.JSON(http.StatusBadRequest, ResponseError{Message: "poolIDs is required"})
+	}
+
+	poolIDs, err := parseNumbers(poolIDsStr)
+	if err != nil {
+		return c.JSON(http.StatusBadRequest, ResponseError{Message: err.Error()})
+	}
+
+	// Quote
+	quote, err := a.RUsecase.GetCustomQuote(ctx, tokenIn, tokenOutDenom, poolIDs)
+	if err != nil {
+		return c.JSON(getStatusCode(err), ResponseError{Message: err.Error()})
+	}
+
+	quote.PrepareResult()
+
+	return c.JSON(http.StatusOK, quote)
+}
+
 // GetCandidateRoutes returns the candidate routes for a given tokenIn and tokenOutDenom
 func (a *RouterHandler) GetCandidateRoutes(c echo.Context) error {
 	ctx := c.Request().Context()
@@ -174,3 +208,31 @@ func getValidTokenInTokenOutStr(c echo.Context) (tokenOutStr, tokenInStr string,
 
 	return tokenOutStr, tokenInStr, nil
 }
+
+// parseNumbers parses a comma-separated list of numbers into a slice of unit64.
+func parseNumbers(numbersParam string) ([]uint64, error) {
+	var numbers []uint64
+	numStrings := splitAndTrim(numbersParam, ",")
+
+	for _, numStr := range numStrings {
+		num, err := strconv.ParseUint(numStr, 10, 64)
+		if err != nil {
+			return nil, err
+		}
+		numbers = append(numbers, num)
+	}
+
+	return numbers, nil
+}
+
+// splitAndTrim splits a string by a separator and trims the resulting strings.
+func splitAndTrim(s, sep string) []string {
+	var result []string
+	for _, val := range strings.Split(s, sep) {
+		trimmed := strings.TrimSpace(val)
+		if trimmed != "" {
+			result = append(result, trimmed)
+		}
+	}
+	return result
+}

ingest/sqs/router/delivery/http/router_handlet_test.go

@@ -0,0 +1,42 @@
+package http_test
+
+import (
+	"reflect"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/osmosis-labs/osmosis/v21/ingest/sqs/router/delivery/http"
+)
+
+// TestParseNumbers tests parsing a string of numbers to a slice of uint64
+func TestParseNumbers(t *testing.T) {
+	testCases := []struct {
+		input           string
+		expectedNumbers []uint64
+		expectedError   bool
+	}{
+		{"", nil, false},                          // Empty string, expecting an empty slice and no error
+		{"1,2,3", []uint64{1, 2, 3}, false},       // Comma-separated numbers, expecting slice {1, 2, 3} and no error
+		{"42", []uint64{42}, false},               // Single number, expecting slice {42} and no error
+		{"10,20,30", []uint64{10, 20, 30}, false}, // Another set of numbers
+
+		// Add more test cases as needed
+		{"abc", nil, true}, // Invalid input, expecting an error
+	}
+
+	for _, testCase := range testCases {
+		actualNumbers, actualError := http.ParseNumbers(testCase.input)
+
+		if testCase.expectedError {
+			require.Error(t, actualError)
+			return
+		}
+
+		// Check if the actual output matches the expected output
+		if !reflect.DeepEqual(actualNumbers, testCase.expectedNumbers) {
+			t.Errorf("Input: %s, Expected Numbers: %v, Actual Numbers: %v",
+				testCase.input, testCase.expectedNumbers, actualNumbers)
+		}
+	}
+}

ingest/sqs/router/usecase/optimized_routes_test.go

@@ -718,6 +718,55 @@ func (s *RouterTestSuite) TestGetBestSplitRoutesQuote_Mainnet_ATOMAKT() {
 	s.Require().NotNil(quote.GetAmountOut())
 }
 
+// Validates custom quote for UOSMO to UION.
+// That is, with the given pool ID, we expect the quote to be routed through the route
+// that matches these pool IDs. Errors otherwise.
+func (s *RouterTestSuite) TestGetCustomQuote_Mainnet_UOSMOUION() {
+	config := defaultRouterConfig
+	config.MaxPoolsPerRoute = 5
+	config.MaxRoutes = 10
+
+	var (
+		amountIn = osmomath.NewInt(5000000)
+	)
+
+	router, tickMap, takerFeeMap := s.setupMainnetRouter(config)
+
+	// Setup router repository mock
+	routerRepositoryMock := mocks.RedisRouterRepositoryMock{
+		TakerFees: takerFeeMap,
+	}
+	routerusecase.WithRouterRepository(router, &routerRepositoryMock)
+
+	// Setup pools usecase mock.
+	poolsRepositoryMock := mocks.RedisPoolsRepositoryMock{
+		Pools:     router.GetSortedPools(),
+		TickModel: tickMap,
+	}
+	poolsUsecase := poolsusecase.NewPoolsUsecase(time.Hour, &poolsRepositoryMock, nil)
+	routerusecase.WithPoolsUsecase(router, poolsUsecase)
+
+	routerUsecase := routerusecase.NewRouterUsecase(time.Hour, &routerRepositoryMock, poolsUsecase, config, &log.NoOpLogger{})
+
+	// This pool ID is second best: https://app.osmosis.zone/pool/2
+	// The top one is https://app.osmosis.zone/pool/1110 which is not selected
+	// due to custom parameter.
+	const expectedPoolID = uint64(2)
+	poolIDs := []uint64{expectedPoolID}
+
+	quote, err := routerUsecase.GetCustomQuote(context.Background(), sdk.NewCoin(UOSMO, amountIn), UION, poolIDs)
+
+	s.Require().NoError(err)
+	s.Require().NotNil(quote)
+
+	s.Require().Len(quote.GetRoute(), 1)
+	routePools := quote.GetRoute()[0].GetPools()
+	s.Require().Len(routePools, 1)
+
+	// Validate that the pool is pool 2
+	s.Require().Equal(expectedPoolID, routePools[0].GetId())
+}
+
 // Generates routes from mainnet state by:
 // - instrumenting pool repository mock with pools and ticks
 // - setting this mock on the pools use case

ingest/sqs/router/usecase/router_usecase.go

@@ -2,6 +2,7 @@ package usecase
 
 import (
 	"context"
+	"fmt"
 	"time"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
@@ -89,6 +90,77 @@ func (r *routerUseCaseImpl) GetBestSingleRouteQuote(ctx context.Context, tokenIn
 	return router.getBestSingleRouteQuote(tokenIn, routes)
 }
 
+// GetCustomQuote implements mvc.RouterUsecase.
+func (r *routerUseCaseImpl) GetCustomQuote(ctx context.Context, tokenIn sdk.Coin, tokenOutDenom string, poolIDs []uint64) (domain.Quote, error) {
+	// TODO: abstract this
+	router := r.initializeRouter()
+
+	candidateRoutes, err := r.handleRoutes(ctx, router, tokenIn.Denom, tokenOutDenom)
+	if err != nil {
+		return nil, err
+	}
+
+	takerFees, err := r.routerRepository.GetAllTakerFees(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	routes, err := r.poolsUsecase.GetRoutesFromCandidates(ctx, candidateRoutes, takerFees, tokenIn.Denom, tokenOutDenom)
+	if err != nil {
+		return nil, err
+	}
+
+	routeIndex := -1
+
+	for curRouteIndex, route := range routes {
+		routePools := route.GetPools()
+
+		// Skip routes that do not match the pool length.
+		if len(routePools) != len(poolIDs) {
+			continue
+		}
+
+		for i, pool := range routePools {
+			poolID := pool.GetId()
+
+			desiredPoolID := poolIDs[i]
+
+			// Break out of the loop if the poolID does not match the desired poolID
+			if poolID != desiredPoolID {
+				break
+			}
+
+			// Found a route that matches the poolIDs
+			if i == len(routePools)-1 {
+				routeIndex = curRouteIndex
+			}
+		}
+
+		// If the routeIndex is not -1, then we found a route that matches the poolIDs
+		// Break out of the loop
+		if routeIndex != -1 {
+			break
+		}
+	}
+
+	// Validate routeIndex
+	if routeIndex == -1 {
+		return nil, fmt.Errorf("no route found for poolIDs: %v", poolIDs)
+	}
+	if routeIndex >= len(routes) {
+		return nil, fmt.Errorf("routeIndex %d is out of bounds", routeIndex)
+	}
+
+	// Compute direct quote
+	foundRoute := routes[routeIndex]
+	quote, _, err := router.estimateBestSingleRouteQuote([]route.RouteImpl{foundRoute}, tokenIn)
+	if err != nil {
+		return nil, err
+	}
+
+	return quote, nil
+}
+
 // GetCandidateRoutes implements domain.RouterUsecase.
 func (r *routerUseCaseImpl) GetCandidateRoutes(ctx context.Context, tokenInDenom string, tokenOutDenom string) (route.CandidateRoutes, error) {
 	router := r.initializeRouter()