Publish a SQLWarning chain as Result Message Segments

Author: Michael-A-McMahon

src/main/java/oracle/r2dbc/impl/OracleResultImpl.java

@@ -34,12 +34,15 @@
 
 import java.sql.BatchUpdateException;
 import java.sql.ResultSet;
+import java.sql.SQLWarning;
+import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.stream.LongStream;
+import java.util.stream.Stream;
 
 import static oracle.r2dbc.impl.OracleR2dbcExceptions.fromJdbc;
 import static oracle.r2dbc.impl.OracleR2dbcExceptions.requireNonNull;
@@ -376,6 +379,20 @@ static OracleResultImpl createErrorResult(R2dbcException r2dbcException) {
     return new ErrorResult(r2dbcException);
   }
 
+  /**
+   * Creates a {@code Result} that publishes a {@code warning} as a
+   * {@link Message} segment, followed by any {@code Segment}s of a
+   * {@code result}.
+   * @param warning Warning to publish
+   * @param result Result to publisher
+   * @return A {@code Result} for a {@code Statement} execution that
+   * completed with a warning.
+   */
+  static OracleResultImpl createWarningResult(
+    SQLWarning warning, OracleResultImpl result) {
+    return new WarningResult(warning, result);
+  }
+
   /**
    * {@link OracleResultImpl} subclass that publishes a single update count. An
    * instance of this class constructed with negative valued update count
@@ -392,16 +409,37 @@ private UpdateCountResult(long updateCount) {
     @Override
     <T> Publisher<T> publishSegments(Function<Segment, T> mappingFunction) {
       return updateCount >= 0
-        ? Mono.just(mappingFunction.apply(new UpdateCountImpl(updateCount)))
+        ? Mono.just(new UpdateCountImpl(updateCount))
+            .map(mappingFunction)
         : Mono.empty();
     }
   }
 
   /**
+   * <p>
    * {@link OracleResultImpl} subclass that publishes JDBC {@link ResultSet} as
    * {@link RowSegment}s. {@link RowMetadata} of published {@code Rows} is
    * derived from the {@link java.sql.ResultSetMetaData} of the
    * {@link ResultSet}.
+   * </p><p>
+   * This {@code Result} is <i>not</i> implemented to publish
+   * {@link SQLWarning} chains returned by {@link ResultSet#getWarnings()} as
+   * {@link Message} segments. This implementation is correct for the 21.1
+   * Oracle JDBC Driver which is known to implement {@code getWarnings()} by
+   * returning {@code null} for forward-only insensitive {@code ResultSets}
+   * when no invocation of {@link java.sql.Statement#setMaxRows(int)}
+   * or {@link java.sql.Statement#setLargeMaxRows(long)} has occurred.
+   * </p><p>
+   * It is a known limitation of the 21.1 Oracle JDBC Driver that
+   * {@link ResultSet#getWarnings()} can not be invoked after row publishing
+   * has been initiated; The {@code ResultSet} is logically closed once row
+   * publishing has been initiated, and so {@code getWarnings} would throw a
+   * {@link java.sql.SQLException} to indicate a closed {@code ResultSet}. If
+   * a later release of Oracle JDBC removes this limitation, then this
+   * {@code Result} should be implemented to invoke {@code getWarnings} to
+   * ensure correctness if a later release of Oracle JDBC also returns non-null
+   * values from that method.
+   * </p>
    */
   private static final class ResultSetResult extends OracleResultImpl {
 
@@ -482,8 +520,8 @@ private BatchUpdateResult(long[] updateCounts) {
     @Override
     <T> Publisher<T> publishSegments(Function<Segment, T> mappingFunction) {
       return Flux.fromStream(LongStream.of(updateCounts)
-        .mapToObj(UpdateCountImpl::new)
-        .map(mappingFunction));
+        .mapToObj(UpdateCountImpl::new))
+        .map(mappingFunction);
     }
   }
 
@@ -532,6 +570,35 @@ <T> Publisher<T> publishSegments(Function<Segment, T> mappingFunction) {
     }
   }
 
+  /**
+   * {@link OracleResultImpl} subclass that publishes a {@link SQLWarning}
+   * chain as {@link Message} segments, followed by the segments of another
+   * {@link OracleResultImpl}.
+   */
+  private static final class WarningResult extends OracleResultImpl {
+
+    private final SQLWarning warning;
+    private final OracleResultImpl result;
+
+    private WarningResult(SQLWarning warning, OracleResultImpl result) {
+      this.warning = warning;
+      this.result = result;
+    }
+
+    @Override
+    <T> Publisher<T> publishSegments(Function<Segment, T> mappingFunction) {
+      return Flux.fromStream(Stream.iterate(
+        warning, Objects::nonNull, SQLWarning::getNextWarning)
+        .map(OracleR2dbcExceptions::toR2dbcException)
+        .map(MessageImpl::new))
+        .map(mappingFunction)
+        // Invoke publishSegments(Class, Function) rather than
+        // publishSegments(Function) to update the state of the result; Namely,
+        // the state that has the onConsumed Publisher emit a terminal signal.
+        .concatWith(result.publishSegments(Segment.class, mappingFunction));
+    }
+  }
+
 
   /**
    * Common interface for instances of {@link Segment} with a {@link Readable}

src/main/java/oracle/r2dbc/impl/OracleStatementImpl.java

@@ -40,6 +40,7 @@
 import java.sql.Connection;
 import java.sql.PreparedStatement;
 import java.sql.SQLType;
+import java.sql.SQLWarning;
 import java.time.Duration;
 import java.util.Arrays;
 import java.util.LinkedList;
@@ -819,7 +820,21 @@ private Publisher<Boolean> publishSqlExecution(
     PreparedStatement preparedStatement, int fetchSize) {
     runJdbc(() -> preparedStatement.setFetchSize(fetchSize));
     setQueryTimeout(preparedStatement);
-    return adapter.publishSQLExecution(preparedStatement);
+    return Mono.from(adapter.publishSQLExecution(preparedStatement))
+      // Work around a bug in 21.1 Oracle JDBC that has the
+      // OraclePreparedStatement.executeAsyncOracle Publisher emit onError
+      // with a SQLException when the database returns a warning. In 21.3 it's
+      // fixed so that the Publisher emits onComplete and the SQLWarning is
+      // obtained from the usual call to PreparedStatement.getWarnings()
+      // TODO: Remove this when 21.1 Oracle JDBC is no longer supported
+      .onErrorResume(error ->
+        // ORA-17110 is the error code for warnings. If the R2dbcException has
+        // this code, then ignore it and return the normal boolean value
+        // indicating if the result is a ResultSet or not
+        error instanceof R2dbcException
+          && ((R2dbcException) error).getErrorCode() == 17110
+          ? Mono.just(null != fromJdbc(preparedStatement::getResultSet))
+          : Mono.error(error));
   }
 
   private void setQueryTimeout(PreparedStatement preparedStatement) {
@@ -1431,6 +1446,7 @@ private static void registerOutParameters(
               .onErrorResume(R2dbcException.class, r2dbcException ->
                 Mono.just(OracleResultImpl.createErrorResult(r2dbcException)))
               .defaultIfEmpty(createUpdateCountResult(-1L))
+              .map(result -> getWarnings(preparedStatement, result))
               .doOnNext(lastResultRef::set),
           discardQueue ->
             Flux.fromIterable(discardQueue)
@@ -1443,6 +1459,25 @@ private static void registerOutParameters(
           .doOnCancel((ThrowingRunnable)preparedStatement::close));
   }
 
+  /**
+   * Returns a {@code Result} that publishes any {@link SQLWarning}s of a
+   * {@code preparedStatement} as {@link io.r2dbc.spi.Result.Message}
+   * segments followed by any {@code Segments} of a {@code result}. This method
+   * returns the provided {@code result} if the {@code preparedStatement} has
+   * no warnings.
+   * @param preparedStatement Statement that may have warnings
+   * @param result Result of executing the {@code preparedStatement}
+   * @return A {@code Result} having any warning messages of the
+   * {@code preparedStatement} along with its execution {@code result}.
+   */
+  private static OracleResultImpl getWarnings(
+    PreparedStatement preparedStatement, OracleResultImpl result) {
+    SQLWarning warning = fromJdbc(preparedStatement::getWarnings);
+    return warning == null
+      ? result
+      : OracleResultImpl.createWarningResult(warning, result);
+  }
+
   /**
    * Returns an exception indicating that a parameter has not been set.
    * @return Unset parameter exception

src/test/java/oracle/r2dbc/impl/OracleStatementImplTest.java

@@ -27,6 +27,8 @@
 import io.r2dbc.spi.R2dbcException;
 import io.r2dbc.spi.R2dbcType;
 import io.r2dbc.spi.Result;
+import io.r2dbc.spi.Result.Message;
+import io.r2dbc.spi.Result.UpdateCount;
 import io.r2dbc.spi.Statement;
 import io.r2dbc.spi.Type;
 import org.junit.jupiter.api.Test;
@@ -36,6 +38,7 @@
 
 import java.math.BigDecimal;
 import java.sql.RowId;
+import java.sql.SQLWarning;
 import java.util.Collections;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Collectors;
@@ -57,7 +60,10 @@
 import static oracle.r2dbc.util.Awaits.tryAwaitExecution;
 import static oracle.r2dbc.util.Awaits.tryAwaitNone;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.fail;
 
 /**
  * Verifies that
@@ -1959,6 +1965,59 @@ public void testOutAndImplicitResult() {
     }
   }
 
+  /**
+   * Verifies that {@link OracleStatementImpl#execute()} emits a {@link Result}
+   * with a {@link Message} segment when the execution results in a
+   * warning.
+   */
+  @Test
+  public void testWarningMessage() {
+    Connection connection =
+      Mono.from(sharedConnection()).block(connectTimeout());
+    try {
+
+      // Create a procedure using invalid syntax and expect the Result to
+      // have a Message with an R2dbcException having a SQLWarning as it's
+      // initial cause. Expect the Result to have an update count of zero as
+      // well, indicating that the statement completed after the warning.
+      AtomicInteger segmentCount = new AtomicInteger(0);
+      R2dbcException r2dbcException =
+        awaitOne(Flux.from(connection.createStatement(
+          "CREATE OR REPLACE PROCEDURE testWarningMessage" +
+            " IS BEGIN;")
+          .execute())
+          .concatMap(result ->
+            result.flatMap(segment -> {
+              int index = segmentCount.getAndIncrement();
+              if (index == 0) {
+                assertTrue(segment instanceof Message,
+                  "Unexpected Segment: " + segment);
+                return Mono.just(((Message)segment).exception());
+              }
+              else if (index == 1) {
+                assertTrue(segment instanceof UpdateCount,
+                  "Unexpected Segment: " + segment);
+                assertEquals(0, ((UpdateCount)segment).value());
+                return Mono.empty();
+              }
+              else {
+                fail("Unexpected Segment: " + segment);
+                return Mono.error(new AssertionError("Should not reach here"));
+              }
+            })));
+
+      // Expect ORA-17110 for an execution that completed with a warning
+      assertEquals(17110, r2dbcException.getErrorCode());
+      Throwable cause = r2dbcException.getCause();
+      assertTrue(cause instanceof SQLWarning, "Unexpected cause: " + cause);
+      assertEquals(17110, ((SQLWarning)cause).getErrorCode());
+      assertNull(cause.getCause());
+    }
+    finally {
+      tryAwaitNone(connection.close());
+    }
+  }
+
   // TODO: Repalce with Parameters.inOut when that's available
   private static final class InOutParameter
     implements Parameter, Parameter.In, Parameter.Out {