Implement docdb pagination

athena-docdb/src/main/java/com/amazonaws/athena/connectors/docdb/DocDBMetadataHandler.java

@@ -43,8 +43,6 @@
 import com.amazonaws.services.glue.model.Table;
 import com.amazonaws.services.secretsmanager.AWSSecretsManager;
 import com.google.common.base.Strings;
-import com.google.common.collect.ImmutableList;
-import com.mongodb.MongoCommandException;
 import com.mongodb.client.MongoClient;
 import com.mongodb.client.MongoCursor;
 import org.apache.arrow.util.VisibleForTesting;
@@ -56,6 +54,10 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+import static com.amazonaws.athena.connector.lambda.metadata.ListTablesRequest.UNLIMITED_PAGE_SIZE_VALUE;
 
 /**
  * Handles metadata requests for the Athena DocumentDB Connector.
@@ -168,22 +170,21 @@ public ListSchemasResponse doListSchemaNames(BlockAllocator blockAllocator, List
     public ListTablesResponse doListTables(BlockAllocator blockAllocator, ListTablesRequest request)
     {
         MongoClient client = getOrCreateConn(request);
+        Stream<String> tableNames = doListTablesWithCommand(client, request);
 
-        try (MongoCursor<String> itr = client.getDatabase(request.getSchemaName()).listCollectionNames().iterator()) {
-            List<TableName> tables = new ArrayList<>();
-            while (itr.hasNext()) {
-                tables.add(new TableName(request.getSchemaName(), itr.next()));
-            }
+        int startToken = request.getNextToken() != null ? Integer.parseInt(request.getNextToken()) : 0;
+        int pageSize = request.getPageSize();
+        String nextToken = null;
 
-            return new ListTablesResponse(request.getCatalogName(), tables, null);
+        if (pageSize != UNLIMITED_PAGE_SIZE_VALUE) {
+            logger.info("Starting at token {} w/ page size {}", startToken, pageSize);
+            tableNames = tableNames.skip(startToken).limit(request.getPageSize());
+            nextToken = Integer.toString(startToken + pageSize);
         }
-        catch (MongoCommandException mongoCommandException) {
-            //do this in failed case instead of replace method in case API changes on doc db.
-            logger.warn("Exception on listCollectionNames on Mongo JAVA client, trying with mongo command line.", mongoCommandException);
-            List<TableName> tableNames = doListTablesWithCommand(client, request);
 
-            return new ListTablesResponse(request.getCatalogName(), tableNames.isEmpty() ? ImmutableList.of() : tableNames, null);
-        }
+        List<TableName> paginatedTables = tableNames.map(tableName -> new TableName(request.getSchemaName(), tableName)).collect(Collectors.toList());
+        logger.info("doListTables returned {} tables. Next token is {}", paginatedTables.size(), nextToken);
+        return new ListTablesResponse(request.getCatalogName(), paginatedTables, nextToken);
     }
 
     /**
@@ -208,18 +209,15 @@ public ListTablesResponse doListTables(BlockAllocator blockAllocator, ListTables
      * @param request
      * @return
      */
-    private List<TableName> doListTablesWithCommand(MongoClient client, ListTablesRequest request)
+    private Stream<String> doListTablesWithCommand(MongoClient client, ListTablesRequest request)
     {
         logger.debug("doListTablesWithCommand Start");
-        List<TableName> tables = new ArrayList<>();
-        Document document = client.getDatabase(request.getSchemaName()).runCommand(new Document("listCollections", 1).append("nameOnly", true).append("authorizedCollections", true));
+        List<String> tables = new ArrayList<>();
+        Document queryDocument = new Document("listCollections", 1).append("nameOnly", true).append("authorizedCollections", true);
+        Document document = client.getDatabase(request.getSchemaName()).runCommand(queryDocument);
 
         List<Document> list = ((Document) document.get("cursor")).getList("firstBatch", Document.class);
-        for (Document doc : list) {
-            tables.add(new TableName(request.getSchemaName(), doc.getString("name")));
-        }
-
-        return tables;
+        return list.stream().map(doc -> doc.getString("name")).sorted();
     }
 
     /**

athena-docdb/src/test/java/com/amazonaws/athena/connectors/docdb/DocDBMetadataHandlerTest.java

@@ -63,6 +63,7 @@
 import org.slf4j.LoggerFactory;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -149,13 +150,21 @@ public void doListTables()
         tableNames.add("table2");
         tableNames.add("table3");
 
+        Document tableNamesDocument = new Document("cursor",
+                new Document("firstBatch",
+                        Arrays.asList(new Document("name", "table1"),
+                                new Document("name", "table2"),
+                                new Document("name", "table3"))));
+
         MongoDatabase mockDatabase = mock(MongoDatabase.class);
         when(mockClient.getDatabase(eq(DEFAULT_SCHEMA))).thenReturn(mockDatabase);
-        when(mockDatabase.listCollectionNames()).thenReturn(StubbingCursor.iterate(tableNames));
+        when(mockDatabase.runCommand(any())).thenReturn(tableNamesDocument);
 
         ListTablesRequest req = new ListTablesRequest(IDENTITY, QUERY_ID, DEFAULT_CATALOG, DEFAULT_SCHEMA,
                 null, UNLIMITED_PAGE_SIZE_VALUE);
+
         ListTablesResponse res = handler.doListTables(allocator, req);
+
         logger.info("doListTables - {}", res.getTables());
 
         for (TableName next : res.getTables()) {

athena-elasticsearch/src/test/java/com/amazonaws/athena/connectors/elasticsearch/ElasticsearchMetadataHandlerTest.java

@@ -56,7 +56,15 @@
 import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
-import java.util.*;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Set;
 
 import static com.amazonaws.athena.connector.lambda.metadata.ListTablesRequest.UNLIMITED_PAGE_SIZE_VALUE;
 import static org.junit.Assert.*;