From 931e49959e8b004bf042736bbfea4954342acfa8 Mon Sep 17 00:00:00 2001
From: Nicholas Molnar <65710+neekolas@users.noreply.github.com>
Date: Thu, 11 Apr 2024 15:38:05 -0700
Subject: [PATCH] Allow larger queries for user preferences (#373)

## tl;dr

We are putting in this hack to allow for larger request sizes for User Preferences requests.

This will be removed once MLS/Groups moves into production and we can start caching these results locally.
---
 pkg/api/message/v1/service.go | 20 +++++++++++++++++---
 pkg/topic/topic.go            | 25 +++++++++++++++----------
 2 files changed, 32 insertions(+), 13 deletions(-)

diff --git a/pkg/api/message/v1/service.go b/pkg/api/message/v1/service.go
index 171a611e..fdf85b87 100644
--- a/pkg/api/message/v1/service.go
+++ b/pkg/api/message/v1/service.go
@@ -38,6 +38,9 @@ const (
 	// maxRowsPerQuery defines the maximum number of rows we can return in a single query
 	maxRowsPerQuery = 100
 
+	// maxUserPreferencesRowsPerQuery sets a higher limit for querying the user preferences table
+	maxUserPreferencesRowsPerQuery = 500
+
 	// maxTopicsPerQueryRequest defines the maximum number of topics that can be queried in a single request.
 	// the number is likely to be more than we want it to be, but would be a safe place to put it -
 	// per Test_LargeQueryTesting, the request decoding already failing before it reaches th handler.
@@ -323,12 +326,13 @@ func (s *Service) SubscribeAll(req *proto.SubscribeAllRequest, stream proto.Mess
 func (s *Service) Query(ctx context.Context, req *proto.QueryRequest) (*proto.QueryResponse, error) {
 	log := s.log.Named("query").With(zap.Strings("content_topics", req.ContentTopics))
 	log.Debug("received request")
+	numContentTopics := len(req.ContentTopics)
 
-	if len(req.ContentTopics) == 0 {
+	if numContentTopics == 0 {
 		return nil, status.Errorf(codes.InvalidArgument, "content topics required")
 	}
 
-	if len(req.ContentTopics) > 1 {
+	if numContentTopics > 1 {
 		// if len(req.ContentTopics) > maxTopicsPerQueryRequest {
 		// 	return nil, status.Errorf(codes.InvalidArgument, "the number of content topics(%d) exceed the maximum topics per query request (%d)", len(req.ContentTopics), maxTopicsPerQueryRequest)
 		// }
@@ -346,7 +350,7 @@ func (s *Service) Query(ctx context.Context, req *proto.QueryRequest) (*proto.Qu
 		}
 	}
 
-	if req.PagingInfo != nil && req.PagingInfo.Limit > maxRowsPerQuery {
+	if req.PagingInfo != nil && int(req.PagingInfo.Limit) > getMaxRows(req.ContentTopics) {
 		return nil, status.Errorf(codes.InvalidArgument, "cannot exceed %d rows per query", maxRowsPerQuery)
 	}
 
@@ -395,3 +399,13 @@ func (s *Service) BatchQuery(ctx context.Context, req *proto.BatchQueryRequest)
 		Responses: responses,
 	}, nil
 }
+
+// Temporarily using this function to allow for flexible limits depending on topic.
+// See: https://github.com/xmtp/xmtp-node-go/pull/373
+func getMaxRows(contentTopics []string) int {
+	if len(contentTopics) == 1 && topic.IsUserPreferences(contentTopics[0]) {
+		return maxUserPreferencesRowsPerQuery
+	}
+
+	return maxRowsPerQuery
+}
diff --git a/pkg/topic/topic.go b/pkg/topic/topic.go
index f2460045..5d7c589c 100644
--- a/pkg/topic/topic.go
+++ b/pkg/topic/topic.go
@@ -5,22 +5,27 @@ import (
 )
 
 var topicCategoryByPrefix = map[string]string{
-	"test":         "test",
-	"contact":      "contact",
-	"intro":        "v1-intro",
-	"dm":           "v1-conversation",
-	"dmE":          "v1-conversation-ephemeral",
-	"invite":       "v2-invite",
-	"groupInvite":  "v2-group-invite",
-	"m":            "v2-conversation",
-	"mE":           "v2-conversation-ephemeral",
-	"privatestore": "private",
+	"test":            "test",
+	"contact":         "contact",
+	"intro":           "v1-intro",
+	"dm":              "v1-conversation",
+	"dmE":             "v1-conversation-ephemeral",
+	"invite":          "v2-invite",
+	"groupInvite":     "v2-group-invite",
+	"m":               "v2-conversation",
+	"mE":              "v2-conversation-ephemeral",
+	"privatestore":    "private",
+	"userpreferences": "userpreferences",
 }
 
 func IsEphemeral(contentTopic string) bool {
 	return Category(contentTopic) == "v2-conversation-ephemeral" || Category(contentTopic) == "v1-conversation-ephemeral"
 }
 
+func IsUserPreferences(contentTopic string) bool {
+	return Category(contentTopic) == "userpreferences"
+}
+
 func Category(contentTopic string) string {
 	if strings.HasPrefix(contentTopic, "test-") {
 		return "test"