[AD-216] Investigate statement.setMaxRows() and see if this needs to be incorporated ins some way during query execution (#66)

### Summary

[AD-216] Investigate statement.setMaxRows() and see if this needs to be incorporated ins some way during query execution

### Description

- Push-down a `$limit` command.

### Related Issue

https://bitquill.atlassian.net/browse/AD-216

* [AD-216] Investigate statement.setMaxRows() and see if this needs to be incorporated ins some way during query execution
https://bitquill.atlassian.net/browse/AD-216
- Push-down a `$limit` command.

* [AD-216] Investigate statement.setMaxRows() and see if this needs to be incorporated ins some way during query execution
https://bitquill.atlassian.net/browse/AD-216
- Fix flaky test.

* Commit Code Coverage Badge

Co-authored-by: birschick-bq <[email protected]>

Author: Bruce Irschick

src/main/java/software/amazon/documentdb/jdbc/DocumentDbQueryExecutor.java

@@ -24,9 +24,10 @@
 import com.mongodb.client.MongoClient;
 import com.mongodb.client.MongoClients;
 import com.mongodb.client.MongoCollection;
-import com.mongodb.client.MongoCursor;
 import com.mongodb.client.MongoDatabase;
+import org.bson.BsonDocument;
 import org.bson.Document;
+import org.bson.conversions.Bson;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import software.amazon.documentdb.jdbc.common.utilities.JdbcColumnMetaData;
@@ -170,21 +171,29 @@ protected java.sql.ResultSet runQuery(final String sql) throws SQLException {
         final MongoDatabase database = client.getDatabase(properties.getDatabase());
         final MongoCollection<Document> collection = database
                 .getCollection(queryContext.getCollectionName());
-        AggregateIterable<Document> iterable = collection
-                .aggregate(queryContext.getAggregateOperations());
+
+        // Only add limit if maxRows is non-zero.
+        final List<Bson> aggregateOperations = queryContext.getAggregateOperations();
+        final long maxRows = statement.getLargeMaxRows();
+        if (maxRows > 0) {
+            // Use push-down to apply maxRows limit.
+            aggregateOperations.add(BsonDocument.parse(
+                    String.format("{\"$limit\": {\"$numberLong\": \"%d\"}}", maxRows)));
+        }
+
+        AggregateIterable<Document> iterable = collection.aggregate(aggregateOperations);
         if (getQueryTimeout() > 0) {
             iterable = iterable.maxTime(getQueryTimeout(), TimeUnit.SECONDS);
         }
         if (getFetchSize() > 0) {
             iterable = iterable.batchSize(getFetchSize());
         }
-        final MongoCursor<Document> iterator = iterable.iterator();
 
         final ImmutableList<JdbcColumnMetaData> columnMetaData = ImmutableList
                 .copyOf(queryContext.getColumnMetaData());
         return new DocumentDbResultSet(
                 this.statement,
-                iterator,
+                iterable.iterator(),
                 columnMetaData,
                 queryContext.getPaths());
     }
@@ -219,7 +228,7 @@ private void performCancel() throws SQLException {
                         SqlError.QUERY_NOT_STARTED_OR_COMPLETE);
             }
 
-            // If there is more than 1 result then more than 1 operations have been given same id
+            // If there is more than 1 result then more than 1 operations have been given same id,
             // and we do not know which to cancel.
             if (ops.size() != 1) {
                 throw SqlError.createSQLException(

src/test/java/software/amazon/documentdb/jdbc/DocumentDbResultSetTest.java

@@ -83,15 +83,15 @@ public class DocumentDbResultSetTest extends DocumentDbFlapDoodleTest {
 
     @BeforeAll
     @SuppressFBWarnings(value = "HARD_CODE_PASSWORD", justification = "Hardcoded for test purposes only")
-    static void initialize() throws SQLException {
+    static void initialize() {
         // Add a valid users to the local MongoDB instance.
         client = createMongoClient("admin", "admin", "admin");
         createUser(DATABASE_NAME, TEST_USER, TEST_PASSWORD);
     }
 
     @BeforeEach
     void init() throws SQLException {
-        MockitoAnnotations.initMocks(this);
+        MockitoAnnotations.openMocks(this);
         Mockito.when(mockStatement.getFetchSize()).thenReturn(MOCK_FETCH_SIZE);
     }
 

src/test/java/software/amazon/documentdb/jdbc/DocumentDbStatementBasicTest.java

@@ -1073,7 +1073,8 @@ void testCurrentTimestamp() throws SQLException {
                 final OffsetDateTime currentDateTime = Instant.now().atOffset(ZoneOffset.UTC);
                 final long diffInMilliSeconds = actualDateTime
                         .until(currentDateTime, ChronoUnit.MILLIS);
-                Assertions.assertTrue(diffInMilliSeconds > 0 && diffInMilliSeconds < 2000);
+                Assertions.assertTrue(diffInMilliSeconds >= 0);
+                Assertions.assertTrue(diffInMilliSeconds < 1000);
                 Assertions.assertFalse(testResultSet.next());
                 return null;
             } catch (SQLException e) {

src/test/java/software/amazon/documentdb/jdbc/DocumentDbStatementFilterTest.java

@@ -22,6 +22,7 @@
 import org.bson.BsonDouble;
 import org.bson.BsonInt64;
 import org.bson.BsonMinKey;
+import org.bson.BsonString;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.DisplayName;
@@ -1399,4 +1400,75 @@ void testDateMinus() throws SQLException {
         Assertions.assertEquals("101", resultSet.getString(1));
         Assertions.assertFalse(resultSet.next());
     }
+
+    @Test
+    @DisplayName("Tests that setMaxRows limits the number of rows returned in result set.")
+    void testSetMaxRows() throws SQLException {
+        final String collection = "testSetMaxRows";
+        final BsonDocument[] documents = new BsonDocument[10];
+        final int totalNumberDocuments = 10;
+        final int maxRows = 5;
+        for (int i = 0; i < totalNumberDocuments; i++) {
+            documents[i] = new BsonDocument("field", new BsonString("value"));
+        }
+        insertBsonDocuments(collection, DATABASE_NAME, USER, PASSWORD, documents);
+        final Statement statement = getDocumentDbStatement();
+
+        // Don't set max rows
+        ResultSet resultSet = statement.executeQuery(
+                String.format("SELECT * FROM \"%s\".\"%s\"", DATABASE_NAME, collection));
+        int actualRowCount = 0;
+        while (resultSet.next()) {
+            actualRowCount++;
+        }
+        Assertions.assertEquals(0, statement.getMaxRows());
+        Assertions.assertEquals(totalNumberDocuments, actualRowCount);
+
+        // Set max rows < actual
+        statement.setMaxRows(maxRows);
+        resultSet = statement.executeQuery(
+                String.format("SELECT * FROM \"%s\".\"%s\"", DATABASE_NAME, collection));
+        actualRowCount = 0;
+        while (resultSet.next()) {
+            actualRowCount++;
+        }
+        Assertions.assertEquals(maxRows, statement.getMaxRows());
+        Assertions.assertEquals(maxRows, actualRowCount);
+
+        // Set unlimited
+        statement.setMaxRows(0);
+        resultSet = statement.executeQuery(
+                String.format("SELECT * FROM \"%s\".\"%s\"", DATABASE_NAME, collection));
+        actualRowCount = 0;
+        while (resultSet.next()) {
+            actualRowCount++;
+        }
+        Assertions.assertEquals(0, statement.getMaxRows());
+        Assertions.assertEquals(totalNumberDocuments, actualRowCount);
+
+        // Set max rows > SQL LIMIT
+        int limit = maxRows - 1;
+        statement.setMaxRows(maxRows);
+        resultSet = statement.executeQuery(
+                String.format("SELECT * FROM \"%s\".\"%s\" LIMIT %d", DATABASE_NAME, collection, limit));
+        actualRowCount = 0;
+        while (resultSet.next()) {
+            actualRowCount++;
+        }
+        Assertions.assertEquals(maxRows, statement.getMaxRows());
+        Assertions.assertEquals(limit, actualRowCount);
+
+        // Set max rows < SQL LIMIT
+        limit = maxRows + 1;
+        statement.setMaxRows(maxRows);
+        resultSet = statement.executeQuery(
+                String.format("SELECT * FROM \"%s\".\"%s\" LIMIT %d", DATABASE_NAME, collection, limit));
+        actualRowCount = 0;
+        while (resultSet.next()) {
+            actualRowCount++;
+        }
+        Assertions.assertEquals(maxRows, statement.getMaxRows());
+        Assertions.assertEquals(maxRows, actualRowCount);
+
+    }
 }