Addressed review comments by using PaginationHelper class.

Author: VenkatasivareddyTR

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

@@ -38,8 +38,6 @@ public class Db2Constants
     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";
-    public static final String ALL_TABLES_COUNT_QUERY = "select count(*) AS \"ALL_TABLES_COUNT\" " +
-            "from syscat.tables where type in ('T', 'U', 'V', 'W') and tabschema = ? ";
 
     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

@@ -45,6 +45,7 @@
 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.jdbc.connection.DatabaseConnectionConfig;
 import com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionInfo;
 import com.amazonaws.athena.connectors.jdbc.connection.GenericJdbcConnectionFactory;
@@ -184,41 +185,21 @@ public ListTablesResponse doListTables(final BlockAllocator blockAllocator, fina
     public ListTablesResponse listPaginatedTables(final Connection connection, final ListTablesRequest listTablesRequest) throws SQLException
     {
         LOGGER.debug("Starting listPaginatedTables for Db2.");
-        String token = listTablesRequest.getNextToken();
         int pageSize = listTablesRequest.getPageSize();
+        int token = PaginationHelper.validateAndParsePaginationArguments(listTablesRequest.getNextToken(), pageSize);
 
         if (pageSize == UNLIMITED_PAGE_SIZE_VALUE) {
-            LOGGER.debug("listPaginatedTables - pagination with UNLIMITED_PAGE_SIZE_VALUE");
-            pageSize = getAllTablesCount(connection, listTablesRequest.getSchemaName());
+            pageSize = Integer.MAX_VALUE;
         }
 
-        int t = token != null ? Integer.parseInt(token) : 0;
+        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);
 
-        LOGGER.info("Starting pagination at {} with page size {}", t, pageSize);
-        List<TableName> paginatedTables = getPaginatedTables(connection, listTablesRequest.getSchemaName(), t, pageSize);
-        LOGGER.info("{} tables returned. Next token is {}", paginatedTables.size(), t + pageSize);
-
-        String nextToken = paginatedTables.isEmpty() || paginatedTables.size() < pageSize || listTablesRequest.getPageSize() == UNLIMITED_PAGE_SIZE_VALUE ? null : Integer.toString(t + pageSize);
-        // return next token is null when reaching end of files
         return new ListTablesResponse(listTablesRequest.getCatalogName(), paginatedTables, nextToken);
     }
 
-    private int getAllTablesCount(Connection connection, String schemaName) throws SQLException
-    {
-        PreparedStatement preparedStatement = connection.prepareStatement(Db2Constants.ALL_TABLES_COUNT_QUERY);
-        preparedStatement.setString(1, schemaName);
-        int allTablesCount = 0;
-        try (ResultSet resultSet = preparedStatement.executeQuery()) {
-            while (resultSet.next()) {
-                allTablesCount = resultSet.getInt(1);
-            }
-        }
-        catch (SQLException sqlException) {
-            throw new RuntimeException(sqlException.getErrorCode() + ": " + sqlException);
-        }
-        return allTablesCount;
-    }
-
     @VisibleForTesting
     protected List<TableName> getPaginatedTables(Connection connection, String databaseName, int offset, int limit) throws SQLException
     {

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

@@ -26,6 +26,7 @@
 import com.amazonaws.athena.connector.lambda.domain.Split;
 import com.amazonaws.athena.connector.lambda.domain.TableName;
 import com.amazonaws.athena.connector.lambda.domain.predicate.Constraints;
+import com.amazonaws.athena.connector.lambda.exceptions.AthenaConnectorException;
 import com.amazonaws.athena.connector.lambda.metadata.GetSplitsRequest;
 import com.amazonaws.athena.connector.lambda.metadata.GetSplitsResponse;
 import com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest;
@@ -401,28 +402,41 @@ blockAllocator, new ListTablesRequest(this.federatedIdentity, "testQueryId",
         Assert.assertNull(listTablesResponse.getNextToken());
         Assert.assertArrayEquals(expected, listTablesResponse.getTables().toArray());
 
-        // Test 4: Testing when requesting pageSize UNLIMITED_PAGE_SIZE_VALUE(-1) and nextToken is 2.
+        // Test 4: nextToken is 2 and pageSize is UNLIMITED. Return all tables starting from index 2.
         preparedStatement = Mockito.mock(PreparedStatement.class);
         Mockito.when(this.connection.prepareStatement(Db2Constants.LIST_PAGINATED_TABLES_QUERY)).thenReturn(preparedStatement);
-        values = new Object[][]{{"testSchema", "testTable2"}};
-        expected = new TableName[]{new TableName("testSchema", "testTable2")};
+        values = new Object[][]{{"testSchema", "testTable3"}, {"testSchema", "testTable4"}};
+        expected = new TableName[]{new TableName("testSchema", "testTable3"), new TableName("testSchema", "testTable4")};
         resultSet = mockResultSet(schema, values, new AtomicInteger(-1));
         Mockito.when(preparedStatement.executeQuery()).thenReturn(resultSet);
 
-        PreparedStatement preparedStatement1 = Mockito.mock(PreparedStatement.class);
-        Mockito.when(this.connection.prepareStatement(Db2Constants.ALL_TABLES_COUNT_QUERY)).thenReturn(preparedStatement1);
-        schema = new String[]{"ALL_TABLES_COUNT"};
-        values = new Object[][] {{5}};
-        resultSet = mockResultSet(schema, values, new AtomicInteger(1));
-        Mockito.when(preparedStatement1.executeQuery()).thenReturn(resultSet);
-        Mockito.when(resultSet.next()).thenReturn(true).thenReturn(false);
-        Mockito.when(resultSet.getInt(1)).thenReturn(5);
-
         listTablesResponse = this.db2MetadataHandler.doListTables(
                 blockAllocator, new ListTablesRequest(this.federatedIdentity, "testQueryId",
                         "testCatalog", "testSchema", "2", ListTablesRequest.UNLIMITED_PAGE_SIZE_VALUE));
         Assert.assertNull(listTablesResponse.getNextToken());
         Assert.assertArrayEquals(expected, listTablesResponse.getTables().toArray());
+
+        // Test 5: AthenaConnectorException with negative nextToken value
+        preparedStatement = Mockito.mock(PreparedStatement.class);
+        Mockito.when(this.connection.prepareStatement(Db2Constants.LIST_PAGINATED_TABLES_QUERY)).thenReturn(preparedStatement);
+        values = new Object[][]{{"testSchema", "testTable3"}, {"testSchema", "testTable4"}};
+        resultSet = mockResultSet(schema, values, new AtomicInteger(-1));
+        Mockito.when(preparedStatement.executeQuery()).thenReturn(resultSet);
+
+        Assert.assertThrows(AthenaConnectorException.class, () -> this.db2MetadataHandler.doListTables(
+                blockAllocator, new ListTablesRequest(this.federatedIdentity, "testQueryId",
+                        "testCatalog", "testSchema", "-1", 3)));
+
+        // Test 6: AthenaConnectorException with negative pageSize value
+        preparedStatement = Mockito.mock(PreparedStatement.class);
+        Mockito.when(this.connection.prepareStatement(Db2Constants.LIST_PAGINATED_TABLES_QUERY)).thenReturn(preparedStatement);
+        values = new Object[][]{{"testSchema", "testTable3"}, {"testSchema", "testTable4"}};
+        resultSet = mockResultSet(schema, values, new AtomicInteger(-1));
+        Mockito.when(preparedStatement.executeQuery()).thenReturn(resultSet);
+
+        Assert.assertThrows(AthenaConnectorException.class, () -> this.db2MetadataHandler.doListTables(
+                blockAllocator, new ListTablesRequest(this.federatedIdentity, "testQueryId",
+                        "testCatalog", "testSchema", "0", -3)));
     }
 }