Added COPY TO STDOUT support

src/main/java/io/r2dbc/postgresql/PostgresqlConnection.java

@@ -16,6 +16,8 @@
 
 package io.r2dbc.postgresql;
 
+import io.netty.util.ReferenceCountUtil;
+import io.netty.util.ReferenceCounted;
 import io.r2dbc.postgresql.api.CopyInBuilder;
 import io.r2dbc.postgresql.api.ErrorDetails;
 import io.r2dbc.postgresql.api.Notification;
@@ -58,6 +60,7 @@
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
 
+import static io.r2dbc.postgresql.PostgresqlStatement.WINDOW_UNTIL;
 import static io.r2dbc.postgresql.client.TransactionStatus.IDLE;
 import static io.r2dbc.postgresql.client.TransactionStatus.OPEN;
 
@@ -233,6 +236,17 @@ public CopyInBuilder copyIn(String sql) {
         return new PostgresqlCopyIn.Builder(this.resources, sql);
     }
 
+    @Override
+    public Flux<io.r2dbc.postgresql.api.PostgresqlCopyOutResult> copyOut(String sql) {
+        ExceptionFactory factory = ExceptionFactory.withSql(sql);
+
+        return SimpleQueryMessageFlow.exchange(this.resources.getClient(), sql)
+            .windowUntil(WINDOW_UNTIL)
+            .doOnDiscard(ReferenceCounted .class, ReferenceCountUtil::release) // ensure release of rows within WindowPredicate
+            .map(messages -> PostgresqlCopyOutResult.toCopyOutResult(messages, factory));
+
+    }
+
     @Override
     public PostgresqlBatch createBatch() {
         return new PostgresqlBatch(this.resources);

src/main/java/io/r2dbc/postgresql/PostgresqlCopyOutResult.java

@@ -0,0 +1,152 @@
+/*
+ * Copyright 2019 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.r2dbc.postgresql;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.util.AbstractReferenceCounted;
+import io.netty.util.ReferenceCountUtil;
+import io.netty.util.ReferenceCounted;
+import io.r2dbc.postgresql.message.Format;
+import io.r2dbc.postgresql.message.backend.*;
+import io.r2dbc.postgresql.util.Assert;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.SynchronousSink;
+
+import java.util.Collection;
+import java.util.Objects;
+import java.util.function.BiFunction;
+
+public class PostgresqlCopyOutResult extends AbstractReferenceCounted implements io.r2dbc.postgresql.api.PostgresqlCopyOutResult {
+
+    private final Flux<BackendMessage> messages;
+
+    private final ExceptionFactory factory;
+
+    private volatile CopyOutMetadata metadata;
+
+    PostgresqlCopyOutResult(Flux<BackendMessage> messages, ExceptionFactory factory) {
+        this.messages = Assert.requireNonNull(messages, "messages must not be null");
+        this.factory = Assert.requireNonNull(factory, "factory must not be null");
+    }
+
+
+    @Override
+    @SuppressWarnings({"rawtypes", "unchecked"})
+    public <T> Flux<T> map(BiFunction<ByteBuf, CopyOutMetadata, ? extends T> f) {
+        Assert.requireNonNull(f, "f must not be null");
+
+        return (Flux<T>) this.messages
+                .handle((message, sink) -> {
+
+                    try {
+                        if (message instanceof ErrorResponse) {
+                            this.factory.handleErrorResponse(message, (SynchronousSink) sink);
+                            return;
+                        }
+
+                        if (message instanceof CopyOutResponse) {
+                            this.metadata = PostgresCopyOutMetadata.toMetadata((CopyOutResponse) message);
+                            return;
+                        }
+
+                        if (message instanceof CopyData) {
+                            CopyData data = (CopyData) message;
+                            sink.next(f.apply(data.getData(), this.metadata));
+                            return;
+                        }
+
+                        if (message instanceof CopyDone) {
+                            sink.complete();
+                            return;
+                        }
+
+                        if (message instanceof RowDescription) {
+                            sink.error(new IllegalStateException("copyOut may only be used with statements that execute a COPY ... TO STDOUT query"));
+                        }
+
+                    } finally {
+                        ReferenceCountUtil.release(message);
+                    }
+                });
+    }
+
+    @Override
+    protected void deallocate() {
+        // drain messages for cleanup
+        messages.map(ReferenceCountUtil::release).subscribe();
+    }
+
+    @Override
+    public ReferenceCounted touch(Object hint) {
+        return this;
+    }
+
+    @Override
+    public String toString() {
+        return "PostgresqlCopyOutResult{" +
+                "messages=" + messages +
+                ", factory=" + factory +
+                '}';
+    }
+
+    static PostgresqlCopyOutResult toCopyOutResult(Flux<BackendMessage> messages, ExceptionFactory factory) {
+        return new PostgresqlCopyOutResult(messages, factory);
+    }
+
+    private static class PostgresCopyOutMetadata implements CopyOutMetadata {
+        private final Format overallFormat;
+        private final Collection<Format> columnFormats;
+
+        private PostgresCopyOutMetadata(Format overallFormat, Collection<Format> columnFormats) {
+            this.overallFormat = overallFormat;
+            this.columnFormats = columnFormats;
+        }
+
+        @Override
+        public Format getOverallFormat() {
+            return overallFormat;
+        }
+
+        public Collection<Format> getColumnFormats() {
+            return columnFormats;
+        }
+
+        @Override
+        public String toString() {
+            return "PostgresCopyOutMetadata{" +
+                    "overallFormat=" + overallFormat +
+                    ", columnFormats=" + columnFormats +
+                    '}';
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (!(o instanceof CopyOutMetadata)) return false;
+            CopyOutMetadata that = (CopyOutMetadata) o;
+            return overallFormat == that.getOverallFormat() && Objects.equals(columnFormats, that.getColumnFormats());
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(overallFormat, columnFormats);
+        }
+
+        static CopyOutMetadata toMetadata(CopyOutResponse response) {
+            return new PostgresCopyOutMetadata(response.getOverallFormat(), response.getColumnFormats());
+        }
+    }
+}

src/main/java/io/r2dbc/postgresql/PostgresqlResult.java

@@ -19,11 +19,7 @@
 import io.netty.util.AbstractReferenceCounted;
 import io.netty.util.ReferenceCountUtil;
 import io.netty.util.ReferenceCounted;
-import io.r2dbc.postgresql.message.backend.BackendMessage;
-import io.r2dbc.postgresql.message.backend.CommandComplete;
-import io.r2dbc.postgresql.message.backend.DataRow;
-import io.r2dbc.postgresql.message.backend.ErrorResponse;
-import io.r2dbc.postgresql.message.backend.RowDescription;
+import io.r2dbc.postgresql.message.backend.*;
 import io.r2dbc.postgresql.util.Assert;
 import io.r2dbc.spi.Result;
 import io.r2dbc.spi.Row;
@@ -127,6 +123,16 @@ public <T> Flux<T> map(BiFunction<Row, RowMetadata, ? extends T> f) {
                     if (message instanceof DataRow) {
                         PostgresqlRow row = PostgresqlRow.toRow(this.resources, (DataRow) message, this.metadata, this.rowDescription);
                         sink.next(f.apply(row, this.metadata));
+                        return;
+                    }
+
+                    if (message instanceof CopyOutResponse) {
+                        sink.error(new IllegalStateException("COPY ... TO STDOUT statements may only be used with PostgresqlConnection.copyOut()"));
+                        return;
+                    }
+
+                    if (message instanceof CopyInResponse) {
+                        sink.error(new IllegalStateException("COPY ... FROM STDIN statements may only be used with PostgresqlConnection.copyIn()"));
                     }
 
                 } finally {

src/main/java/io/r2dbc/postgresql/PostgresqlStatement.java

@@ -57,7 +57,7 @@
  */
 final class PostgresqlStatement implements io.r2dbc.postgresql.api.PostgresqlStatement {
 
-    private static final Predicate<BackendMessage> WINDOW_UNTIL = or(CommandComplete.class::isInstance, EmptyQueryResponse.class::isInstance, ErrorResponse.class::isInstance);
+    static final Predicate<BackendMessage> WINDOW_UNTIL = or(CommandComplete.class::isInstance, EmptyQueryResponse.class::isInstance, ErrorResponse.class::isInstance);
 
     private final ArrayDeque<Binding> bindings;
 

src/main/java/io/r2dbc/postgresql/api/PostgresqlConnection.java

@@ -18,11 +18,7 @@
 
 import io.netty.buffer.ByteBuf;
 import io.r2dbc.postgresql.message.frontend.CancelRequest;
-import io.r2dbc.spi.Connection;
-import io.r2dbc.spi.IsolationLevel;
-import io.r2dbc.spi.R2dbcNonTransientResourceException;
-import io.r2dbc.spi.TransactionDefinition;
-import io.r2dbc.spi.ValidationDepth;
+import io.r2dbc.spi.*;
 import org.reactivestreams.Publisher;
 import org.reactivestreams.Subscriber;
 import reactor.core.publisher.Flux;
@@ -70,6 +66,15 @@ public interface PostgresqlConnection extends Connection {
      */
     CopyInBuilder copyIn(String sql);
 
+    /**
+     * Executes one or more PostgreSQL COPY ... TO STDOUT statements and returns the {@link PostgresqlCopyOutResult}s.
+     * {@link PostgresqlCopyOutResult} objects must be fully consumed to ensure full execution of the {@link Statement}.
+     *
+     * @return the {@link PostgresqlCopyOutResult}s, returned by each statement
+     * @since 1.2
+     */
+    Flux<io.r2dbc.postgresql.api.PostgresqlCopyOutResult> copyOut(String sql);
+
     /**
      * Use {@code COPY FROM STDIN} for very fast copying into a database table.
      *

src/main/java/io/r2dbc/postgresql/api/PostgresqlCopyOutResult.java

@@ -0,0 +1,70 @@
+/*
+ * Copyright 2019 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.r2dbc.postgresql.api;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.util.ReferenceCounted;
+import io.r2dbc.postgresql.message.Format;
+import reactor.core.publisher.Flux;
+
+import java.util.Collection;
+import java.util.function.BiFunction;
+
+/**
+ * A result obtained by executing a {@code COPY ... TO STDOUT} statement.
+ */
+public interface PostgresqlCopyOutResult {
+
+    /**
+     * Returns a mapping of the rows that are the results of a copy-out operation against a database. May be empty if
+     * the query did not return any rows.
+     * <p>According to the PostgreSQL documentation, each {@link ByteBuf} will represent a complete row either in
+     * binary or textual representation, however it is not clear whether that is just an incidental implementation
+     * quirk or a strict guarantee. In the other direction, there is no such requirement, so robust consumers should
+     * probably not make any assumptions here.
+     * <p>The {@link ByteBuf} instances are only valid to use during the {@code mappingFunction} callback, unless they
+     * are retained (via {@link ReferenceCounted#retain()} and later released. When retaining them, care must be taken
+     * to ensure that they do eventually get released, as the asynchronous nature of streams can often see messages
+     * dropped when the stream is canceled downstream.
+     *
+     * @param mappingFunction the {@link BiFunction} that maps a {@link ByteBuf} and {@link CopyOutMetadata} to a value
+     * @param <T>             the type of the mapped value
+     * @return a mapping of the rows that are the results of a copy-out operation against a database
+     * @throws IllegalArgumentException if {@code mappingFunction} is {@code null}
+     * @throws IllegalStateException    if the result was consumed
+     */
+    <T> Flux<T> map(BiFunction<ByteBuf, CopyOutMetadata, ? extends T> mappingFunction);
+
+    /**
+     * Metadata for a copy-out operation.
+     */
+    interface CopyOutMetadata {
+        /**
+         * Get the format of each of the columns.
+         *
+         * @return The format of each of the columns
+         */
+        Collection<Format> getColumnFormats();
+
+        /**
+         * Get the overall format of a response.
+         *
+         * @return The format.
+         */
+        Format getOverallFormat();
+    }
+}