From 34680c5ccf1578a1c8531532265c3ea218be329c Mon Sep 17 00:00:00 2001
From: Harshavardhana <harsha@minio.io>
Date: Tue, 2 Nov 2021 19:14:46 -0700
Subject: [PATCH] fix: SQL select to honor limits properly for array queries
 (#13568)

added tests to cover the scenarios as well.
---
 internal/s3select/select.go          | 11 +++++
 internal/s3select/select_test.go     | 68 ++++++++++++++++++++++++++++
 internal/s3select/sql/parser_test.go |  1 +
 internal/s3select/sql/statement.go   |  9 ++--
 4 files changed, 83 insertions(+), 6 deletions(-)

diff --git a/internal/s3select/select.go b/internal/s3select/select.go
index dbdcaeecc..b292dd6d4 100644
--- a/internal/s3select/select.go
+++ b/internal/s3select/select.go
@@ -548,6 +548,17 @@ OuterLoop:
 				}
 
 				outputQueue[len(outputQueue)-1] = outputRecord
+				if s3Select.statement.LimitReached() {
+					if !sendRecord() {
+						break
+					}
+					if err = writer.Finish(s3Select.getProgress()); err != nil {
+						// FIXME: log this error.
+						err = nil
+					}
+					return
+				}
+
 				if len(outputQueue) < cap(outputQueue) {
 					continue
 				}
diff --git a/internal/s3select/select_test.go b/internal/s3select/select_test.go
index d383c169e..bee0314cb 100644
--- a/internal/s3select/select_test.go
+++ b/internal/s3select/select_test.go
@@ -404,6 +404,74 @@ func TestJSONQueries(t *testing.T) {
 			query:      `SELECT 3.0 / 2, 5 / 2.0 FROM S3Object LIMIT 1`,
 			wantResult: `{"_1":1.5,"_2":2.5}`,
 		},
+		{
+			name:  "limit-1",
+			query: `SELECT * FROM S3Object[*].elements[*] LIMIT 1`,
+			requestXML: []byte(`
+<?xml version="1.0" encoding="UTF-8"?>
+<SelectObjectContentRequest>
+    <Expression>select * from s3object[*].elements[*] s where s.element_type = '__elem__merfu'</Expression>
+    <ExpressionType>SQL</ExpressionType>
+    <InputSerialization>
+        <CompressionType>NONE</CompressionType>
+        <JSON>
+            <Type>DOCUMENT</Type>
+        </JSON>
+    </InputSerialization>
+    <OutputSerialization>
+        <JSON>
+        </JSON>
+    </OutputSerialization>
+    <RequestProgress>
+        <Enabled>FALSE</Enabled>
+    </RequestProgress>
+</SelectObjectContentRequest>`),
+			wantResult: `{"element_type":"__elem__merfu","element_id":"d868aefe-ef9a-4be2-b9b2-c9fd89cc43eb","attributes":{"__attr__image_dpi":300,"__attr__image_size":[2550,3299],"__attr__image_index":2,"__attr__image_format":"JPEG","__attr__file_extension":"jpg","__attr__data":null}}`,
+			withJSON: `
+{
+  "name": "small_pdf1.pdf",
+  "lume_id": "9507193e-572d-4f95-bcf1-e9226d96be65",
+  "elements": [
+    {
+      "element_type": "__elem__image",
+      "element_id": "859d09c4-7cf1-4a37-9674-3a7de8b56abc",
+      "attributes": {
+        "__attr__image_dpi": 300,
+        "__attr__image_size": [
+          2550,
+          3299
+        ],
+        "__attr__image_index": 1,
+        "__attr__image_format": "JPEG",
+        "__attr__file_extension": "jpg",
+        "__attr__data": null
+      }
+    },
+    {
+      "element_type": "__elem__merfu",
+      "element_id": "d868aefe-ef9a-4be2-b9b2-c9fd89cc43eb",
+      "attributes": {
+        "__attr__image_dpi": 300,
+        "__attr__image_size": [
+          2550,
+          3299
+        ],
+        "__attr__image_index": 2,
+        "__attr__image_format": "JPEG",
+        "__attr__file_extension": "jpg",
+        "__attr__data": null
+      }
+    }
+  ],
+  "data": "asdascasdc1234e123erdasdas"
+}`,
+		},
+		{
+			name:       "limit-2",
+			query:      `select * from s3object[*].person[*] limit 1`,
+			wantResult: `{"Id":1,"Name":"Anshu","Address":"Templestowe","Car":"Jeep"}`,
+			withJSON:   `{ "person": [ { "Id": 1, "Name": "Anshu", "Address": "Templestowe", "Car": "Jeep" }, { "Id": 2, "Name": "Ben Mostafa", "Address": "Las Vegas", "Car": "Mustang" }, { "Id": 3, "Name": "Rohan Wood", "Address": "Wooddon", "Car": "VW" } ] }`,
+		},
 	}
 
 	defRequest := `<?xml version="1.0" encoding="UTF-8"?>
diff --git a/internal/s3select/sql/parser_test.go b/internal/s3select/sql/parser_test.go
index bc402553c..ff0ed3a84 100644
--- a/internal/s3select/sql/parser_test.go
+++ b/internal/s3select/sql/parser_test.go
@@ -332,6 +332,7 @@ func TestFromClauseJSONPath(t *testing.T) {
 		"select * from s3object[*].books[*].name as s",
 		"select * from s3object s where name > 2",
 		"select * from s3object[*].name as s where name > 2",
+		"select * from s3object[*].books[*] limit 1",
 	}
 	for i, tc := range cases {
 		err := p.ParseString(tc, &s)
diff --git a/internal/s3select/sql/statement.go b/internal/s3select/sql/statement.go
index 6ed7df385..0f11384cc 100644
--- a/internal/s3select/sql/statement.go
+++ b/internal/s3select/sql/statement.go
@@ -302,9 +302,8 @@ func (e *SelectStatement) Eval(input, output Record) (Record, error) {
 		// .. WHERE ..`
 
 		// Update count of records output.
-		if e.limitValue > -1 {
-			e.outputCount++
-		}
+		e.outputCount++
+
 		return input.Clone(output), nil
 	}
 
@@ -328,9 +327,7 @@ func (e *SelectStatement) Eval(input, output Record) (Record, error) {
 	}
 
 	// Update count of records output.
-	if e.limitValue > -1 {
-		e.outputCount++
-	}
+	e.outputCount++
 
 	return output, nil
 }