From c5970b2cdfd7981c1e9c2e3cf440e5f438bb499d Mon Sep 17 00:00:00 2001
From: Dan Hansen <dan@recidiviz.org>
Date: Thu, 25 Jan 2024 13:11:53 -0800
Subject: [PATCH 1/3] [Queries] Respsect `maxResults` parameter when fetching
 query results

---
 internal/types/types.go |  2 +-
 server/handler.go       | 29 +++++++++++++++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/internal/types/types.go b/internal/types/types.go
index ba8bef6cc..fe718e63f 100644
--- a/internal/types/types.go
+++ b/internal/types/types.go
@@ -14,7 +14,7 @@ type (
 	GetQueryResultsResponse struct {
 		JobReference *bigqueryv2.JobReference `json:"jobReference"`
 		Schema       *bigqueryv2.TableSchema  `json:"schema"`
-		Rows         []*TableRow              `json:"rows"`
+		Rows         []*TableRow              `json:"rows,omitempty"`
 		TotalRows    uint64                   `json:"totalRows,string"`
 		JobComplete  bool                     `json:"jobComplete"`
 		TotalBytes   uint64                   `json:"-"`
diff --git a/server/handler.go b/server/handler.go
index 022efaf10..77a43a383 100644
--- a/server/handler.go
+++ b/server/handler.go
@@ -555,6 +555,8 @@ func (h *uploadContentHandler) Handle(ctx context.Context, r *uploadContentReque
 const (
 	formatOptionsUseInt64TimestampParam = "formatOptions.useInt64Timestamp"
 	deleteContentsParam                 = "deleteContents"
+	maxResults                          = "maxResults"
+	maxResultsDefaultValue              = -1
 )
 
 func isDeleteContents(r *http.Request) bool {
@@ -581,6 +583,22 @@ func parseQueryValueAsBool(r *http.Request, key string) bool {
 	return b
 }
 
+func parseQueryValueAsInt64(r *http.Request, key string, defaultValue int64) int64 {
+	queryValues := r.URL.Query()
+	values, exists := queryValues[key]
+	if !exists {
+		return defaultValue
+	}
+	if len(values) != 1 {
+		return defaultValue
+	}
+	parsed, err := strconv.ParseInt(values[0], 10, 64)
+	if err != nil {
+		return defaultValue
+	}
+	return parsed
+}
+
 func (h *datasetsDeleteHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	ctx := r.Context()
 	server := serverFromContext(ctx)
@@ -966,6 +984,7 @@ func (h *jobsGetQueryResultsHandler) ServeHTTP(w http.ResponseWriter, r *http.Re
 		server:            server,
 		project:           project,
 		job:               job,
+		maxResults:        parseQueryValueAsInt64(r, maxResults, maxResultsDefaultValue),
 		useInt64Timestamp: isFormatOptionsUseInt64Timestamp(r),
 	})
 	if err != nil {
@@ -980,6 +999,7 @@ type jobsGetQueryResultsRequest struct {
 	project           *metadata.Project
 	job               *metadata.Job
 	useInt64Timestamp bool
+	maxResults        int64
 }
 
 func (h *jobsGetQueryResultsHandler) Handle(ctx context.Context, r *jobsGetQueryResultsRequest) (*internaltypes.GetQueryResultsResponse, error) {
@@ -988,6 +1008,15 @@ func (h *jobsGetQueryResultsHandler) Handle(ctx context.Context, r *jobsGetQuery
 		return nil, err
 	}
 	rows := internaltypes.Format(response.Schema, response.Rows, r.useInt64Timestamp)
+
+	if r.maxResults == 0 {
+		rows = nil
+	}
+
+	if r.maxResults != maxResultsDefaultValue {
+		rows = rows[:r.maxResults]
+	}
+
 	return &internaltypes.GetQueryResultsResponse{
 		JobReference: &bigqueryv2.JobReference{
 			ProjectId: r.project.ID,

From 05ade4eecbf796ee286886ec5774c718fbbc0964 Mon Sep 17 00:00:00 2001
From: Dan Hansen <dan@recidiviz.org>
Date: Thu, 25 Jan 2024 15:07:21 -0800
Subject: [PATCH 2/3] limit slice to max bounds

---
 server/handler.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/server/handler.go b/server/handler.go
index 77a43a383..9f4440538 100644
--- a/server/handler.go
+++ b/server/handler.go
@@ -1014,7 +1014,7 @@ func (h *jobsGetQueryResultsHandler) Handle(ctx context.Context, r *jobsGetQuery
 	}
 
 	if r.maxResults != maxResultsDefaultValue {
-		rows = rows[:r.maxResults]
+		rows = rows[:min(int64(len(rows)), r.maxResults)]
 	}
 
 	return &internaltypes.GetQueryResultsResponse{

From efb0d46ea2170e3374ac98a1d376d382ad40258b Mon Sep 17 00:00:00 2001
From: Dan Hansen <dan@recidiviz.org>
Date: Mon, 8 Apr 2024 09:29:33 -0700
Subject: [PATCH 3/3] validate maxResults

---
 server/handler.go | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/server/handler.go b/server/handler.go
index 9f4440538..a32805ea4 100644
--- a/server/handler.go
+++ b/server/handler.go
@@ -1013,6 +1013,10 @@ func (h *jobsGetQueryResultsHandler) Handle(ctx context.Context, r *jobsGetQuery
 		rows = nil
 	}
 
+	if r.maxResults < -1 {
+		return nil, fmt.Errorf("invalid maxResults parameter; must be greater than or equal to [-1], got [%i]", r.maxResults)
+	}
+
 	if r.maxResults != maxResultsDefaultValue {
 		rows = rows[:min(int64(len(rows)), r.maxResults)]
 	}