Resolves #3037: Only allow continuation to be returned after a result…

… set has been exhausted (#3038)

* Only allow continuations for exhausted result sets

* Change tests to match new continuation logic

* Added comment

* PR Comments

fdb-relational-api/src/main/java/com/apple/foundationdb/relational/api/RelationalResultSet.java

@@ -55,6 +55,7 @@ public interface RelationalResultSet extends java.sql.ResultSet, RelationalStruc
 
     /**
      * A {@code Continuation} that can be used for retrieving the rest of the rows.
+     * Note: The continuation is only available once the result set has been exhausted.
      *
      * @return  A {@code Continuation} that can be used for retrieving the rest of the rows.
      * @throws SQLException if the continuation cannot be retrieved.

fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/IteratorResultSet.java

@@ -21,18 +21,18 @@
 package com.apple.foundationdb.relational.recordlayer;
 
 import com.apple.foundationdb.annotation.API;
-
 import com.apple.foundationdb.relational.api.Continuation;
 import com.apple.foundationdb.relational.api.Row;
 import com.apple.foundationdb.relational.api.StructMetaData;
 import com.apple.foundationdb.relational.api.exceptions.RelationalException;
 import com.apple.foundationdb.relational.util.SpotBugsSuppressWarnings;
 
 import javax.annotation.Nonnull;
-import java.nio.ByteBuffer;
 import java.sql.SQLException;
 import java.util.Iterator;
 
+import static com.apple.foundationdb.relational.api.exceptions.ErrorCode.UNSUPPORTED_OPERATION;
+
 
 /**
  * Placeholder result set until better generic abstractions come along.
@@ -56,15 +56,13 @@ public IteratorResultSet(StructMetaData metaData, Iterator<? extends Row> rowIte
     public Continuation getContinuation() throws SQLException {
         boolean hasNext = rowIter.hasNext();
         boolean beginning = currentRowPosition == 0;
-        int currPos = currentRowPosition + 1;
+
+        if (hasNext) {
+            throw new SQLException("Continuation can only be returned once the result set has been exhausted", UNSUPPORTED_OPERATION.getErrorCode());
+        }
 
         if (beginning) {
             return ContinuationImpl.BEGIN;
-        } else if (hasNext) {
-            ByteBuffer buffer = ByteBuffer.allocate(4);
-            buffer.putInt(currPos);
-            buffer.flip();
-            return ContinuationImpl.fromUnderlyingBytes(buffer.array());
         } else {
             return ContinuationImpl.END;
         }

fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/RecordLayerResultSet.java

@@ -33,6 +33,8 @@
 import javax.annotation.Nullable;
 import java.sql.SQLException;
 
+import static com.apple.foundationdb.relational.api.exceptions.ErrorCode.UNSUPPORTED_OPERATION;
+
 @API(API.Status.EXPERIMENTAL)
 public class RecordLayerResultSet extends AbstractRecordLayerResultSet {
 
@@ -106,6 +108,9 @@ public boolean isClosed() throws SQLException {
     @Nonnull
     @Override
     public Continuation getContinuation() throws SQLException {
+        if (hasNext()) {
+            throw new SQLException("Continuation can only be returned once the result set has been exhausted", UNSUPPORTED_OPERATION.getErrorCode());
+        }
         try {
             return enrichContinuationFunction.apply(currentCursor.getContinuation(), continuationReason());
         } catch (RelationalException e) {

fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/CursorTest.java

@@ -21,7 +21,6 @@
 package com.apple.foundationdb.relational.recordlayer;
 
 import com.apple.foundationdb.relational.api.Continuation;
-import com.apple.foundationdb.relational.api.ImmutableRowStruct;
 import com.apple.foundationdb.relational.api.KeySet;
 import com.apple.foundationdb.relational.api.Options;
 import com.apple.foundationdb.relational.api.Row;
@@ -36,7 +35,6 @@
 import com.apple.foundationdb.relational.utils.ResultSetTestUtils;
 import com.apple.foundationdb.relational.utils.SimpleDatabaseRule;
 import com.apple.foundationdb.relational.utils.TestSchemas;
-import com.apple.foundationdb.relational.utils.RelationalStructAssert;
 
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Order;
@@ -104,24 +102,22 @@ public void canIterateWithContinuation() throws SQLException, RelationalExceptio
 
     @Test
     public void continuationOnEdgesOfRecordCollection() throws SQLException, RelationalException {
-        insertRecordsAndTest(3, (List<RelationalStruct> records, RelationalConnection conn) -> {
+        final int numRecords = 3;
+        insertRecordsAndTest(numRecords, (List<RelationalStruct> records, RelationalConnection conn) -> {
             try (RelationalResultSet resultSet = conn.createStatement().executeScan("RESTAURANT", new KeySet(), Options.NONE)) {
-                // get continuation before iterating on the result set (should point to the first record).
-                Continuation continuation = resultSet.getContinuation();
-                Assertions.assertEquals(ContinuationImpl.BEGIN, continuation, "Incorrect starting continuation!");
+                // get continuation before iterating on the result set (should fail).
+                Assertions.assertThrows(SQLException.class, resultSet::getContinuation);
 
                 StructMetaData smd = resultSet.getMetaData();
-                boolean called = false;
+                int called = 0;
                 while (resultSet.next()) {
-                    called = true;
-                    Row resumedRow = readFirstRecordWithContinuation(conn.createStatement(), continuation);
-                    Row mainRow = ResultSetTestUtils.currentRow(resultSet);
-                    RelationalStructAssert.assertThat(new ImmutableRowStruct(mainRow, smd)).isEqualTo(new ImmutableRowStruct(resumedRow, smd));
-
-                    continuation = resultSet.getContinuation();
+                    // Continuation is not available until the result set is exhausted
+                    if (++called < numRecords) {
+                        Assertions.assertThrows(SQLException.class, resultSet::getContinuation);
+                    }
                 }
 
-                Assertions.assertTrue(called, "Did not return any records!");
+                Assertions.assertTrue(called > 0, "Did not return any records!");
 
                 // get continuation at the last record (should point to FINISHED).
                 Continuation lastContinuation = resultSet.getContinuation();
@@ -130,7 +126,7 @@ public void continuationOnEdgesOfRecordCollection() throws SQLException, Relatio
                 Assertions.assertTrue(lastContinuation.atEnd());
                 Assertions.assertEquals(lastContinuation.getExecutionState(), ContinuationImpl.END.getExecutionState());
 
-            } catch (RelationalException | SQLException e) {
+            } catch (SQLException e) {
                 Assertions.fail("failed to parse ", e);
             }
         });
@@ -169,7 +165,7 @@ public void continuationWithReturnRowLimit() throws SQLException, RelationalExce
             try (final var s = conn.createStatement()) {
                 s.setMaxRows(5);
                 try (final var resultSet = s.executeQuery("select * from RESTAURANT")) {
-                    Assertions.assertTrue(resultSet.getContinuation().atBeginning());
+                    Assertions.assertThrows(SQLException.class, resultSet::getContinuation);
                     final var resultSetAssert = ResultSetAssert.assertThat(resultSet);
                     for (int i = 0; i < 5; i++) {
                         resultSetAssert.hasNextRow();
@@ -183,7 +179,7 @@ public void continuationWithReturnRowLimit() throws SQLException, RelationalExce
             try (final var preparedStatement = conn.prepareStatement("select * from RESTAURANT with continuation ?param")) {
                 preparedStatement.setBytes("param", continuation.serialize());
                 try (final var resultSet = preparedStatement.executeQuery()) {
-                    Assertions.assertTrue(resultSet.getContinuation().atBeginning());
+                    Assertions.assertThrows(SQLException.class, resultSet::getContinuation);
                     final var resultSetAssert = ResultSetAssert.assertThat(resultSet);
                     for (int i = 0; i < 5; i++) {
                         resultSetAssert.hasNextRow();
@@ -207,7 +203,7 @@ public void continuationWithScanRowLimit() throws SQLException, RelationalExcept
         try (final var conn = driver.connect(database.getConnectionUri(), Options.builder().withOption(Options.Name.EXECUTION_SCANNED_ROWS_LIMIT, 3).build())) {
             conn.setSchema(database.getSchemaName());
             try (final var resultSet = conn.createStatement().executeQuery("select * from RESTAURANT")) {
-                Assertions.assertTrue(resultSet.getContinuation().atBeginning());
+                Assertions.assertThrows(SQLException.class, resultSet::getContinuation);
                 while (true) {
                     if (resultSet.next()) {
                         numRowsReturned++;
@@ -225,7 +221,7 @@ public void continuationWithScanRowLimit() throws SQLException, RelationalExcept
             try (final var preparedStatement = conn.prepareStatement("select * from RESTAURANT with continuation ?param")) {
                 preparedStatement.setBytes("param", continuation.serialize());
                 try (final var resultSet = preparedStatement.executeQuery()) {
-                    Assertions.assertTrue(resultSet.getContinuation().atBeginning());
+                    Assertions.assertThrows(SQLException.class, resultSet::getContinuation);
                     while (true) {
                         if (resultSet.next()) {
                             numRowsReturned++;

fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/IteratorResultSetTest.java

@@ -42,9 +42,7 @@ void continuationAtBeginning() throws RelationalException, SQLException {
                 FieldDescription.primitive("testField", Types.VARCHAR, DatabaseMetaData.columnNoNulls)
         };
         try (IteratorResultSet irs = new IteratorResultSet(new RelationalStructMetaData(fields), Collections.singleton((Row) (new ArrayRow(new Object[]{"test"}))).iterator(), 0)) {
-            Continuation shouldBeStart = irs.getContinuation();
-            Assertions.assertTrue(shouldBeStart.atBeginning(), "Is not at beginning!");
-            Assertions.assertNull(shouldBeStart.getExecutionState(), "Incorrect byte[] for continuation!");
+            Assertions.assertThrows(SQLException.class, () -> irs.getContinuation());
 
             //now iterate
             irs.next();

fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/JoinWithLimitTest.java

@@ -110,7 +110,7 @@ void innerCorrelatedJoinLimit1() throws Exception {
     @Test
     void innerCorrelatedJoinWithContinuationAndLimit() throws Exception {
         Continuation continuation;
-        statement.setMaxRows(3);
+        statement.setMaxRows(1);
         try (var resultSet = statement.executeQuery("SELECT M.e FROM (SELECT * FROM Q, Q.d) as M")) {
             Assertions.assertThat(resultSet.next()).isTrue();
             Assertions.assertThat(resultSet.getLong("e")).isEqualTo(100L);
@@ -235,7 +235,7 @@ void joinWithLimitLargerThanTableSize() throws Exception {
     @Test
     void joinWithContinuationAndLimit() throws Exception {
         Continuation continuation;
-        statement.setMaxRows(2);
+        statement.setMaxRows(1);
         try (var resultSet = statement.executeQuery("SELECT rpk, sa FROM R, S")) {
             Assertions.assertThat(resultSet.next()).isTrue();
             Assertions.assertThat(resultSet.getLong("rpk")).isEqualTo(1L);

fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/catalog/RecordLayerStoreCatalogTestBase.java

@@ -102,8 +102,8 @@ void testListSchemas() throws RelationalException, SQLException {
             Continuation continuation = ContinuationImpl.BEGIN;
             do {
                 try (RelationalResultSet result = storeCatalog.listSchemas(listTxn, continuation)) {
-                    // to test continuation, only read 1 result at once
-                    if (result.next()) {
+                    // to test continuation, read all results and then continue with the continuation
+                    while (result.next()) {
                         fullSchemaNames.add(result.getString("DATABASE_ID") + "?schema=" + result.getString("SCHEMA_NAME"));
                     }
                     continuation = result.getContinuation();
@@ -122,7 +122,7 @@ void testListSchemas() throws RelationalException, SQLException {
             Continuation continuation = ContinuationImpl.BEGIN;
             do {
                 try (RelationalResultSet result = storeCatalog.listSchemas(listTxn, URI.create("/TEST/test_database_id1"), continuation)) {
-                    if (result.next()) {
+                    while (result.next()) {
                         resultSet.add(result.getString("DATABASE_ID") + "?schema=" + result.getString("SCHEMA_NAME"));
                     }
                     continuation = result.getContinuation();

fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/StandardQueryTests.java

@@ -369,6 +369,7 @@ void selectWithFalsePredicate2() throws Exception {
     void selectWithContinuation() throws Exception {
         try (var ddl = Ddl.builder().database(URI.create("/TEST/QT")).relationalExtension(relationalExtension).schemaTemplate(schemaTemplate).build()) {
             try (var statement = ddl.setSchemaAndGetConnection().createStatement()) {
+                statement.setMaxRows(1);
                 insertRestaurantComplexRecord(statement);
                 RelationalStruct l42 = insertRestaurantComplexRecord(statement, 42L, "rest1");
                 RelationalStruct l43 = insertRestaurantComplexRecord(statement, 43L, "rest1");
@@ -387,8 +388,12 @@ void selectWithContinuation() throws Exception {
                     try (final RelationalResultSet resultSet = statement.executeQuery(query)) {
                         // assert result matches expected
                         Assertions.assertNotNull(resultSet, "Did not return a result set!");
-                        ResultSetAssert.assertThat(resultSet).hasNextRow()
-                                .isRowPartly(expected.get(i));
+                        if (i < expected.size()) {
+                            ResultSetAssert.assertThat(resultSet).hasNextRow()
+                                    .isRowPartly(expected.get(i));
+                        } else {
+                            ResultSetAssert.assertThat(resultSet).hasNoNextRow();
+                        }
                         // get continuation for the next query
                         continuation = resultSet.getContinuation();
                         i += 1;