Added pagination for Db2 connector (awslabs#2772)

Co-authored-by: burhan94 <[email protected]>
Co-authored-by: Fal Bharadwaj <[email protected]>

Author: 3 people

athena-db2/src/main/java/com/amazonaws/athena/connectors/db2/Db2Constants.java

@@ -35,6 +35,9 @@ public class Db2Constants
             "from syscat.tables " +
             "where type in ('T', 'U', 'V', 'W') and " +
             "tabschema = ? order by tabname";
+    public static final String LIST_PAGINATED_TABLES_QUERY = "select tabschema AS \"TABLE_SCHEM\", tabname AS \"TABLE_NAME\" " +
+            "from syscat.tables where type in ('T', 'U', 'V', 'W') and tabschema = ? " +
+            "ORDER BY tabname OFFSET ? ROWS FETCH NEXT ? ROWS ONLY";
 
     static final String PARTITION_QUERY = "SELECT DATAPARTITIONID FROM SYSCAT.DATAPARTITIONS WHERE TABSCHEMA = ? AND TABNAME = ? AND SEQNO > 0";
     static final String COLUMN_INFO_QUERY = "select colname, typename from syscat.columns where tabschema = ? AND tabname = ?";

athena-db2/src/main/java/com/amazonaws/athena/connectors/db2/Db2MetadataHandler.java

@@ -37,12 +37,15 @@
 import com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest;
 import com.amazonaws.athena.connector.lambda.metadata.ListSchemasRequest;
 import com.amazonaws.athena.connector.lambda.metadata.ListSchemasResponse;
+import com.amazonaws.athena.connector.lambda.metadata.ListTablesRequest;
+import com.amazonaws.athena.connector.lambda.metadata.ListTablesResponse;
 import com.amazonaws.athena.connector.lambda.metadata.optimizations.DataSourceOptimizations;
 import com.amazonaws.athena.connector.lambda.metadata.optimizations.OptimizationSubType;
 import com.amazonaws.athena.connector.lambda.metadata.optimizations.pushdown.ComplexExpressionPushdownSubType;
 import com.amazonaws.athena.connector.lambda.metadata.optimizations.pushdown.FilterPushdownSubType;
 import com.amazonaws.athena.connector.lambda.metadata.optimizations.pushdown.LimitPushdownSubType;
 import com.amazonaws.athena.connector.lambda.metadata.optimizations.pushdown.TopNPushdownSubType;
+import com.amazonaws.athena.connector.util.PaginationHelper;
 import com.amazonaws.athena.connectors.db2.resolver.Db2JDBCCaseResolver;
 import com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig;
 import com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionInfo;
@@ -80,6 +83,7 @@
 import java.util.stream.Collectors;
 
 import static com.amazonaws.athena.connector.lambda.domain.predicate.functions.StandardFunctions.NULLIF_FUNCTION_NAME;
+import static com.amazonaws.athena.connector.lambda.metadata.ListTablesRequest.UNLIMITED_PAGE_SIZE_VALUE;
 import static com.amazonaws.athena.connectors.db2.Db2Constants.PARTITION_NUMBER;
 
 public class Db2MetadataHandler extends JdbcMetadataHandler
@@ -170,6 +174,37 @@ protected List<TableName> listTables(Connection connection, String schemaName) t
                 .collect(Collectors.toList());
     }
 
+    @Override
+    public ListTablesResponse listPaginatedTables(final Connection connection, final ListTablesRequest listTablesRequest) throws SQLException
+    {
+        LOGGER.debug("Starting listPaginatedTables for Db2.");
+        int pageSize = listTablesRequest.getPageSize();
+        int token = PaginationHelper.validateAndParsePaginationArguments(listTablesRequest.getNextToken(), pageSize);
+
+        if (pageSize == UNLIMITED_PAGE_SIZE_VALUE) {
+            pageSize = Integer.MAX_VALUE;
+        }
+
+        LOGGER.info("Starting pagination at {} with page size {}", token, pageSize);
+        List<TableName> paginatedTables = getPaginatedTables(connection, listTablesRequest.getSchemaName(), token, pageSize);
+        String nextToken = PaginationHelper.calculateNextToken(token, pageSize, paginatedTables);
+        LOGGER.info("{} tables returned. Next token is {}", paginatedTables.size(), nextToken);
+
+        return new ListTablesResponse(listTablesRequest.getCatalogName(), paginatedTables, nextToken);
+    }
+
+    @VisibleForTesting
+    protected List<TableName> getPaginatedTables(Connection connection, String databaseName, int offset, int limit) throws SQLException
+    {
+        PreparedStatement preparedStatement = connection.prepareStatement(Db2Constants.LIST_PAGINATED_TABLES_QUERY);
+
+        preparedStatement.setString(1, databaseName);
+        preparedStatement.setInt(2, offset);
+        preparedStatement.setInt(3, limit);
+
+        return JDBCUtil.getTableMetadata(preparedStatement, TABLES_AND_VIEWS);
+    }
+
     /**
      * Creates Schema object with arrow compatible filed to frame the partition.
      *

athena-db2/src/test/java/com/amazonaws/athena/connectors/db2/Db2MetadataHandlerTest.java

@@ -72,7 +72,6 @@
 
 import static com.amazonaws.athena.connector.lambda.metadata.ListTablesRequest.UNLIMITED_PAGE_SIZE_VALUE;
 import static com.amazonaws.athena.connectors.db2.Db2Constants.PARTITION_NUMBER;
-import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.nullable;
 
 public class Db2MetadataHandlerTest extends TestBase {
@@ -352,5 +351,54 @@ public void doListTables() throws Exception {
                 new TableName("TESTSCHEMA", "testTABLE")};
         Assert.assertEquals(Arrays.toString(expectedTables), listTablesResponse.getTables().toString());
     }
+
+    @Test
+    public void doListPaginatedTables()
+            throws Exception
+    {
+        //Test 1: Testing Single table returned in request of page size 1 and nextToken null
+        Object[][] values = {{"testSchema", "testTable"}};
+        TableName[] expected = {new TableName("testSchema", "testTable")};
+        executePaginatedTableTest(values, expected, null, 1, "1");
+
+        // Test 2: Testing next table returned of page size 1 and nextToken 1
+        executePaginatedTableTest(values, expected, "1", 1, "2");
+
+        // Test 3: Testing single table returned when requesting pageSize 2 signifying end of pagination where nextToken is null.
+        executePaginatedTableTest(values, expected, "2", 2, null);
+
+        // Test 4: Testing unlimited page size (UNLIMITED_PAGE_SIZE_VALUE) which should set pageSize to Integer.MAX_VALUE
+        values = new Object[][]{{"testSchema", "testTable1"}, {"testSchema", "testTable2"}};
+        expected = new TableName[]{new TableName("testSchema", "testTable1"), new TableName("testSchema", "testTable2")};
+        // Use a non-null token to force pagination path, which will trigger the listPaginatedTables method
+        // With unlimited page size, there should be no next token (null) since all results are returned
+        executePaginatedTableTest(values, expected, "0", UNLIMITED_PAGE_SIZE_VALUE, null);
+    }
+
+    /**
+     * Helper method to execute paginated table test and verify results
+     * @param values the table data to return in the ResultSet
+     * @param expected the expected TableName array
+     * @param nextToken the nextToken to use in the request
+     * @param pageSize the pageSize to use in the request
+     * @param expectedNextToken the expected nextToken in the response
+     * @throws Exception if there's an error during execution
+     */
+    private void executePaginatedTableTest(Object[][] values,
+                                           TableName[] expected, String nextToken, int pageSize,
+                                           String expectedNextToken) throws Exception {
+        BlockAllocator blockAllocator = new BlockAllocatorImpl();
+        PreparedStatement preparedStatement = Mockito.mock(PreparedStatement.class);
+        Mockito.when(this.connection.prepareStatement(Db2Constants.LIST_PAGINATED_TABLES_QUERY)).thenReturn(preparedStatement);
+        String[] schema = {"TABLE_SCHEM", "TABLE_NAME"};
+        ResultSet resultSet = mockResultSet(schema, values, new AtomicInteger(-1));
+        Mockito.when(preparedStatement.executeQuery()).thenReturn(resultSet);
+        ListTablesResponse listTablesResponse = this.db2MetadataHandler.doListTables(
+                blockAllocator, new ListTablesRequest(this.federatedIdentity, "testQueryId",
+                        "testCatalog", "testSchema", nextToken, pageSize));
+
+        Assert.assertEquals(expectedNextToken, listTablesResponse.getNextToken());
+        Assert.assertArrayEquals(expected, listTablesResponse.getTables().toArray());
+    }
 }