From 866ca99baa1a6cccd0703cf74439b338629ae654 Mon Sep 17 00:00:00 2001
From: Nicholas Molnar <65710+neekolas@users.noreply.github.com>
Date: Thu, 11 Apr 2024 13:24:53 -0700
Subject: [PATCH] Allow larger queries for user preferences

---
 pkg/api/message/v1/service.go | 13 ++++++++++++-
 pkg/topic/topic.go            | 25 +++++++++++++++----------
 2 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/pkg/api/message/v1/service.go b/pkg/api/message/v1/service.go
index 171a611e..6390661d 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.
@@ -346,7 +349,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[0]) {
 		return nil, status.Errorf(codes.InvalidArgument, "cannot exceed %d rows per query", maxRowsPerQuery)
 	}
 
@@ -395,3 +398,11 @@ func (s *Service) BatchQuery(ctx context.Context, req *proto.BatchQueryRequest)
 		Responses: responses,
 	}, nil
 }
+
+func getMaxRows(contentTopic string) int {
+	if topic.IsUserPreferences(contentTopic) {
+		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"