Add initial support for result segment mapping.

mp911de · mp911de · commit 9efe389e4bbf · 2021-07-19T11:33:08.000+02:00
Signed-off-by: Mark Paluch <mpaluch@vmware.com> [#199]
diff --git a/src/main/java/io/r2dbc/mssql/ExceptionFactory.java b/src/main/java/io/r2dbc/mssql/ExceptionFactory.java
@@ -90,20 +90,20 @@ static R2dbcException createException(AbstractInfoToken token, String sql) {
             case 565: // A stack overflow occurred in the server while compiling the query. Please simplify the query.
             case 4408: // Too many tables. The query and the views or functions in it exceed the limit of %d tables. Revise the query to reduce the number of tables.
             case 2812: // Could not find stored procedure '%.*ls'.
-                return new MssqlBadGrammarException(createExceptionDetails(token), sql);
+                return new MssqlBadGrammarException(createErrorDetails(token), sql);
 
             case 2601: // Cannot insert duplicate key row in object '%.*ls' with unique index '%.*ls'. The duplicate key value is %ls.
             case 2627: // Violation of %ls constraint '%.*ls'. Cannot insert duplicate key in object '%.*ls'. The duplicate key value is %ls.
             case 544: // Cannot insert explicit value for identity column in table '%.*ls' when IDENTITY_INSERT is set to OFF.
             case 8114: // Error converting data type %ls to %ls.
             case 8115:  // Arithmetic overflow error converting %ls to data type %ls.
-                return new MssqlDataIntegrityViolationException(createExceptionDetails(token));
+                return new MssqlDataIntegrityViolationException(createErrorDetails(token));
 
             case 701: // Maximum number of databases used for each query has been exceeded. The maximum allowed is %d.
             case 1222: // Lock request time out period exceeded.
             case 1204: // The instance of the SQL Server Database Engine cannot obtain a LOCK resource at this time. Rerun your statement when there are fewer active users. Ask the database
                 // administrator to check the lock and memory configuration for this instance, or to check for
-                return new MssqlTransientException(createExceptionDetails(token));
+                return new MssqlTransientException(createErrorDetails(token));
 
             case 1203: // Process ID %d attempted to unlock a resource it does not own: %.*ls. Retry the transaction, because this error may be caused by a timing condition. If the problem
                 // persists, contact the database administrator.
@@ -117,7 +117,7 @@ static R2dbcException createException(AbstractInfoToken token, String sql) {
             case 3938: // The transaction has been stopped because it conflicted with the execution of a FILESTREAM close operation using the same transaction.  The transaction will be rolled back.
             case 28611: // The request is aborted because the transaction has been aborted by Matrix Transaction Coordination Manager. This is mostly caused by one or more transaction particpant
                 // brick went offline.
-                return new MssqlRollbackException(createExceptionDetails(token));
+                return new MssqlRollbackException(createErrorDetails(token));
 
             case 921:  // Database '%.*ls' has not been recovered yet. Wait and try again.
             case 941:  // Database '%.*ls' cannot be opened because it is not started. Retry when the database is started.
@@ -131,7 +131,7 @@ static R2dbcException createException(AbstractInfoToken token, String sql) {
             case 40642: // The server is currently too busy.  Please try again later.
             case 40675: // The service is currently too busy.  Please try again later.
             case 40825: // Unable to complete request now. Please try again later.
-                return new MssqlTransientResourceException(createExceptionDetails(token));
+                return new MssqlTransientResourceException(createErrorDetails(token));
         }
 
         if (token.getClassification() == GENERAL_ERROR && token.getNumber() == 4002) {
@@ -141,19 +141,19 @@ static R2dbcException createException(AbstractInfoToken token, String sql) {
         switch (token.getClassification()) {
             case OBJECT_DOES_NOT_EXIST:
             case SYNTAX_ERROR:
-                return new MssqlBadGrammarException(createExceptionDetails(token), sql);
+                return new MssqlBadGrammarException(createErrorDetails(token), sql);
             case INCONSISTENT_NO_LOCK:
-                return new MssqlDataIntegrityViolationException(createExceptionDetails(token));
+                return new MssqlDataIntegrityViolationException(createErrorDetails(token));
             case TX_DEADLOCK:
-                return new MssqlRollbackException(createExceptionDetails(token));
+                return new MssqlRollbackException(createErrorDetails(token));
             case SECURITY:
-                return new MssqlPermissionDeniedException(createExceptionDetails(token));
+                return new MssqlPermissionDeniedException(createErrorDetails(token));
             case GENERAL_ERROR:
-                return new MssqlNonTransientException(createExceptionDetails(token));
+                return new MssqlNonTransientException(createErrorDetails(token));
             case OUT_OF_RESOURCES:
-                return new MssqlTransientResourceException(createExceptionDetails(token));
+                return new MssqlTransientResourceException(createErrorDetails(token));
             default:
-                return new MssqlNonTransientResourceException(createExceptionDetails(token));
+                return new MssqlNonTransientResourceException(createErrorDetails(token));
         }
     }
 
@@ -181,7 +181,7 @@ RuntimeException createException(ErrorToken message) {
         return createException(message, this.sql);
     }
 
-    private static ErrorDetails createExceptionDetails(AbstractInfoToken token) {
+    static ErrorDetails createErrorDetails(AbstractInfoToken token) {
         return new ErrorDetails(token.getMessage(), token.getNumber(), token.getState(), token.getInfoClass(), token.getServerName(), token.getProcName(), token.getLineNumber());
     }
 
diff --git a/src/main/java/io/r2dbc/mssql/MssqlResult.java b/src/main/java/io/r2dbc/mssql/MssqlResult.java
@@ -19,8 +19,8 @@
 import io.netty.util.ReferenceCountUtil;
 import io.r2dbc.mssql.client.ConnectionContext;
 import io.r2dbc.mssql.codec.Codecs;
-import io.r2dbc.mssql.message.Message;
 import io.r2dbc.mssql.message.token.AbstractDoneToken;
+import io.r2dbc.mssql.message.token.AbstractInfoToken;
 import io.r2dbc.mssql.message.token.ColumnMetadataToken;
 import io.r2dbc.mssql.message.token.ErrorToken;
 import io.r2dbc.mssql.message.token.NbcRowToken;
@@ -30,12 +30,15 @@
 import io.r2dbc.spi.Result;
 import io.r2dbc.spi.Row;
 import io.r2dbc.spi.RowMetadata;
+import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.util.Logger;
 import reactor.util.Loggers;
 
 import java.util.function.BiFunction;
+import java.util.function.Function;
+import java.util.function.Predicate;
 
 /**
  * {@link Result} of query results.
@@ -54,13 +57,13 @@ public final class MssqlResult implements Result {
 
     private final Codecs codecs;
 
-    private final Flux<Message> messages;
+    private final Flux<io.r2dbc.mssql.message.Message> messages;
 
     private volatile MssqlRowMetadata rowMetadata;
 
     private volatile RuntimeException throwable;
 
-    private MssqlResult(String sql, ConnectionContext context, Codecs codecs, Flux<Message> messages) {
+    private MssqlResult(String sql, ConnectionContext context, Codecs codecs, Flux<io.r2dbc.mssql.message.Message> messages) {
 
         this.sql = sql;
         this.context = context;
@@ -76,7 +79,7 @@ private MssqlResult(String sql, ConnectionContext context, Codecs codecs, Flux<M
      * @param messages message stream.
      * @return {@link Result} object.
      */
-    static MssqlResult toResult(String sql, ConnectionContext context, Codecs codecs, Flux<Message> messages) {
+    static MssqlResult toResult(String sql, ConnectionContext context, Codecs codecs, Flux<io.r2dbc.mssql.message.Message> messages) {
 
         Assert.requireNonNull(sql, "SQL must not be null");
         Assert.requireNonNull(codecs, "Codecs must not be null");
@@ -197,4 +200,151 @@ public <T> Flux<T> map(BiFunction<Row, RowMetadata, ? extends T> f) {
             });
     }
 
+    @Override
+    public Result filter(Predicate<Segment> filter) {
+
+        Flux<io.r2dbc.mssql.message.Message> filteredMessages = this.messages.filter(message -> {
+
+            if (message.getClass() == ColumnMetadataToken.class) {
+
+                ColumnMetadataToken token = (ColumnMetadataToken) message;
+
+                if (token.hasColumns()) {
+                    this.rowMetadata = MssqlRowMetadata.create(this.codecs, token);
+                }
+                return true;
+            }
+
+            if (message.getClass() == RowToken.class || message.getClass() == NbcRowToken.class) {
+
+                MssqlRowMetadata rowMetadata = this.rowMetadata;
+
+                if (rowMetadata == null) {
+                    return false;
+                }
+
+                MssqlRow row = MssqlRow.toRow(this.codecs, (RowToken) message, rowMetadata);
+
+                boolean result = filter.test(row);
+
+                if (!result) {
+                    row.release();
+                }
+
+                return result;
+            }
+
+            if (message instanceof AbstractInfoToken) {
+                return filter.test(createMessage((AbstractInfoToken) message));
+            }
+
+            if (message instanceof AbstractDoneToken) {
+
+                AbstractDoneToken doneToken = (AbstractDoneToken) message;
+                if (doneToken.hasCount()) {
+
+                    return filter.test(doneToken);
+                }
+            }
+            return true;
+        });
+
+        return new MssqlResult(this.sql, this.context, this.codecs, filteredMessages);
+    }
+
+    @Override
+    public <T> Flux<T> flatMap(Function<Segment, ? extends Publisher<? extends T>> mappingFunction) {
+
+        return this.messages
+            .flatMap(message -> {
+
+                if (message instanceof AbstractDoneToken) {
+
+                    AbstractDoneToken doneToken = (AbstractDoneToken) message;
+                    if (doneToken.hasCount()) {
+
+                        if (DEBUG_ENABLED) {
+                            LOGGER.debug(this.context.getMessage("Incoming row count: {}"), doneToken);
+                        }
+
+                        return mappingFunction.apply(doneToken);
+                    }
+                }
+
+                if (message.getClass() == ColumnMetadataToken.class) {
+
+                    ColumnMetadataToken token = (ColumnMetadataToken) message;
+
+                    if (!token.hasColumns()) {
+                        return Mono.empty();
+                    }
+
+                    if (DEBUG_ENABLED) {
+                        LOGGER.debug(this.context.getMessage("Result column definition: {}"), message);
+                    }
+
+                    this.rowMetadata = MssqlRowMetadata.create(this.codecs, token);
+                }
+
+                if (message.getClass() == RowToken.class || message.getClass() == NbcRowToken.class) {
+
+                    MssqlRowMetadata rowMetadata1 = this.rowMetadata;
+
+                    if (rowMetadata1 == null) {
+                        return Mono.error(new IllegalStateException("No MssqlRowMetadata available"));
+                    }
+
+                    MssqlRow row = MssqlRow.toRow(this.codecs, (RowToken) message, rowMetadata1);
+
+                    try {
+                        return Flux.from(mappingFunction.apply(row)).doFinally(it -> row.release());
+                    } catch (RuntimeException e) {
+                        row.release();
+                        throw e;
+                    }
+                }
+
+                if (message instanceof AbstractInfoToken) {
+                    return mappingFunction.apply(createMessage((AbstractInfoToken) message));
+                }
+
+                ReferenceCountUtil.release(message);
+
+                return Mono.empty();
+            });
+    }
+
+    private Message createMessage(AbstractInfoToken message) {
+
+        ErrorDetails errorDetails = ExceptionFactory.createErrorDetails(message);
+
+        return new Message() {
+
+            @Override
+            public R2dbcException exception() {
+                return ExceptionFactory.createException(message, MssqlResult.this.sql);
+            }
+
+            @Override
+            public int errorCode() {
+                return (int) errorDetails.getNumber();
+            }
+
+            @Override
+            public String sqlState() {
+                return errorDetails.getStateCode();
+            }
+
+            @Override
+            public String message() {
+                return errorDetails.getMessage();
+            }
+
+            @Override
+            public Severity severity() {
+                return message instanceof ErrorToken ? Severity.ERROR : Severity.INFO;
+            }
+        };
+    }
+
 }
diff --git a/src/main/java/io/r2dbc/mssql/MssqlRow.java b/src/main/java/io/r2dbc/mssql/MssqlRow.java
@@ -23,7 +23,9 @@
 import io.r2dbc.mssql.message.token.RowToken;
 import io.r2dbc.mssql.message.type.SqlServerType;
 import io.r2dbc.mssql.util.Assert;
+import io.r2dbc.spi.Result;
 import io.r2dbc.spi.Row;
+import io.r2dbc.spi.RowMetadata;
 import reactor.util.annotation.Nullable;
 
 /**
@@ -35,7 +37,7 @@
  * @see #release()
  * @see ReferenceCounted
  */
-final class MssqlRow implements Row {
+final class MssqlRow implements Row, Result.Data {
 
     private static final int STATE_ACTIVE = 0;
 
@@ -82,6 +84,11 @@ public MssqlRowMetadata getMetadata() {
         return this.metadata;
     }
 
+    @Override
+    public RowMetadata metadata() {
+        return this.metadata;
+    }
+
     @Override
     public <T> T get(int index, Class<T> type) {
 
diff --git a/src/main/java/io/r2dbc/mssql/message/token/AbstractDoneToken.java b/src/main/java/io/r2dbc/mssql/message/token/AbstractDoneToken.java
@@ -20,6 +20,7 @@
 import io.r2dbc.mssql.message.Message;
 import io.r2dbc.mssql.message.tds.Encode;
 import io.r2dbc.mssql.util.Assert;
+import io.r2dbc.spi.Result;
 
 import java.util.Objects;
 
@@ -28,7 +29,7 @@
  *
  * @author Mark Paluch
  */
-public abstract class AbstractDoneToken extends AbstractDataToken {
+public abstract class AbstractDoneToken extends AbstractDataToken implements Result.UpdateCount {
 
     /**
      * Packet length in bytes.
@@ -246,4 +247,9 @@ public String toString() {
         return sb.toString();
     }
 
+    @Override
+    public long value() {
+        return hasCount() ? getRowCount() : 0;
+    }
+
 }
diff --git a/src/test/java/io/r2dbc/mssql/SimpleMssqlStatementUnitTests.java b/src/test/java/io/r2dbc/mssql/SimpleMssqlStatementUnitTests.java
