Skip to content

Commit a0d55a3

Browse files
jroperjroper
committed
Added COPY TO STDOUT support
[resolves #332] This adds a new PostgresqlConnection.copyOut method to support COPY ... TO STDOUT statements. In addition, PostgresqlResult now fails when executing these statements, rather than silently dropping all the data.
1 parent 1874ffd commit a0d55a3

9 files changed

Lines changed: 450 additions & 11 deletions

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@
1616

1717
package io.r2dbc.postgresql;
1818

19+
import io.netty.util.ReferenceCountUtil;
20+
import io.netty.util.ReferenceCounted;
1921
import io.r2dbc.postgresql.api.CopyInBuilder;
2022
import io.r2dbc.postgresql.api.ErrorDetails;
2123
import io.r2dbc.postgresql.api.Notification;
@@ -58,6 +60,7 @@
5860
import java.util.concurrent.atomic.AtomicReference;
5961
import java.util.function.Function;
6062

63+
import static io.r2dbc.postgresql.PostgresqlStatement.WINDOW_UNTIL;
6164
import static io.r2dbc.postgresql.client.TransactionStatus.IDLE;
6265
import static io.r2dbc.postgresql.client.TransactionStatus.OPEN;
6366

@@ -233,6 +236,17 @@ public CopyInBuilder copyIn(String sql) {
233236
return new PostgresqlCopyIn.Builder(this.resources, sql);
234237
}
235238

239+
@Override
240+
public Flux<io.r2dbc.postgresql.api.PostgresqlCopyOutResult> copyOut(String sql) {
241+
ExceptionFactory factory = ExceptionFactory.withSql(sql);
242+
243+
return SimpleQueryMessageFlow.exchange(this.resources.getClient(), sql)
244+
.windowUntil(WINDOW_UNTIL)
245+
.doOnDiscard(ReferenceCounted .class, ReferenceCountUtil::release) // ensure release of rows within WindowPredicate
246+
.map(messages -> PostgresqlCopyOutResult.toCopyOutResult(messages, factory));
247+
248+
}
249+
236250
@Override
237251
public PostgresqlBatch createBatch() {
238252
return new PostgresqlBatch(this.resources);

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

Lines changed: 152 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,152 @@
1+
/*
2+
* Copyright 2019 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package io.r2dbc.postgresql;
18+
19+
import io.netty.buffer.ByteBuf;
20+
import io.netty.util.AbstractReferenceCounted;
21+
import io.netty.util.ReferenceCountUtil;
22+
import io.netty.util.ReferenceCounted;
23+
import io.r2dbc.postgresql.message.Format;
24+
import io.r2dbc.postgresql.message.backend.*;
25+
import io.r2dbc.postgresql.util.Assert;
26+
import reactor.core.publisher.Flux;
27+
import reactor.core.publisher.SynchronousSink;
28+
29+
import java.util.Collection;
30+
import java.util.Objects;
31+
import java.util.function.BiFunction;
32+
33+
public class PostgresqlCopyOutResult extends AbstractReferenceCounted implements io.r2dbc.postgresql.api.PostgresqlCopyOutResult {
34+
35+
private final Flux<BackendMessage> messages;
36+
37+
private final ExceptionFactory factory;
38+
39+
private volatile CopyOutMetadata metadata;
40+
41+
PostgresqlCopyOutResult(Flux<BackendMessage> messages, ExceptionFactory factory) {
42+
this.messages = Assert.requireNonNull(messages, "messages must not be null");
43+
this.factory = Assert.requireNonNull(factory, "factory must not be null");
44+
}
45+
46+
47+
@Override
48+
@SuppressWarnings({"rawtypes", "unchecked"})
49+
public <T> Flux<T> map(BiFunction<ByteBuf, CopyOutMetadata, ? extends T> f) {
50+
Assert.requireNonNull(f, "f must not be null");
51+
52+
return (Flux<T>) this.messages
53+
.handle((message, sink) -> {
54+
55+
try {
56+
if (message instanceof ErrorResponse) {
57+
this.factory.handleErrorResponse(message, (SynchronousSink) sink);
58+
return;
59+
}
60+
61+
if (message instanceof CopyOutResponse) {
62+
this.metadata = PostgresCopyOutMetadata.toMetadata((CopyOutResponse) message);
63+
return;
64+
}
65+
66+
if (message instanceof CopyData) {
67+
CopyData data = (CopyData) message;
68+
sink.next(f.apply(data.getData(), this.metadata));
69+
return;
70+
}
71+
72+
if (message instanceof CopyDone) {
73+
sink.complete();
74+
return;
75+
}
76+
77+
if (message instanceof RowDescription) {
78+
sink.error(new IllegalStateException("copyOut may only be used with statements that execute a COPY ... TO STDOUT query"));
79+
}
80+
81+
} finally {
82+
ReferenceCountUtil.release(message);
83+
}
84+
});
85+
}
86+
87+
@Override
88+
protected void deallocate() {
89+
// drain messages for cleanup
90+
messages.map(ReferenceCountUtil::release).subscribe();
91+
}
92+
93+
@Override
94+
public ReferenceCounted touch(Object hint) {
95+
return this;
96+
}
97+
98+
@Override
99+
public String toString() {
100+
return "PostgresqlCopyOutResult{" +
101+
"messages=" + messages +
102+
", factory=" + factory +
103+
'}';
104+
}
105+
106+
static PostgresqlCopyOutResult toCopyOutResult(Flux<BackendMessage> messages, ExceptionFactory factory) {
107+
return new PostgresqlCopyOutResult(messages, factory);
108+
}
109+
110+
private static class PostgresCopyOutMetadata implements CopyOutMetadata {
111+
private final Format overallFormat;
112+
private final Collection<Format> columnFormats;
113+
114+
private PostgresCopyOutMetadata(Format overallFormat, Collection<Format> columnFormats) {
115+
this.overallFormat = overallFormat;
116+
this.columnFormats = columnFormats;
117+
}
118+
119+
@Override
120+
public Format getOverallFormat() {
121+
return overallFormat;
122+
}
123+
124+
public Collection<Format> getColumnFormats() {
125+
return columnFormats;
126+
}
127+
128+
@Override
129+
public String toString() {
130+
return "PostgresCopyOutMetadata{" +
131+
"overallFormat=" + overallFormat +
132+
", columnFormats=" + columnFormats +
133+
'}';
134+
}
135+
136+
@Override
137+
public boolean equals(Object o) {
138+
if (!(o instanceof CopyOutMetadata)) return false;
139+
CopyOutMetadata that = (CopyOutMetadata) o;
140+
return overallFormat == that.getOverallFormat() && Objects.equals(columnFormats, that.getColumnFormats());
141+
}
142+
143+
@Override
144+
public int hashCode() {
145+
return Objects.hash(overallFormat, columnFormats);
146+
}
147+
148+
static CopyOutMetadata toMetadata(CopyOutResponse response) {
149+
return new PostgresCopyOutMetadata(response.getOverallFormat(), response.getColumnFormats());
150+
}
151+
}
152+
}

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

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,7 @@
1919
import io.netty.util.AbstractReferenceCounted;
2020
import io.netty.util.ReferenceCountUtil;
2121
import io.netty.util.ReferenceCounted;
22-
import io.r2dbc.postgresql.message.backend.BackendMessage;
23-
import io.r2dbc.postgresql.message.backend.CommandComplete;
24-
import io.r2dbc.postgresql.message.backend.DataRow;
25-
import io.r2dbc.postgresql.message.backend.ErrorResponse;
26-
import io.r2dbc.postgresql.message.backend.RowDescription;
22+
import io.r2dbc.postgresql.message.backend.*;
2723
import io.r2dbc.postgresql.util.Assert;
2824
import io.r2dbc.spi.Result;
2925
import io.r2dbc.spi.Row;
@@ -127,6 +123,16 @@ public <T> Flux<T> map(BiFunction<Row, RowMetadata, ? extends T> f) {
127123
if (message instanceof DataRow) {
128124
PostgresqlRow row = PostgresqlRow.toRow(this.resources, (DataRow) message, this.metadata, this.rowDescription);
129125
sink.next(f.apply(row, this.metadata));
126+
return;
127+
}
128+
129+
if (message instanceof CopyOutResponse) {
130+
sink.error(new IllegalStateException("COPY ... TO STDOUT statements may only be used with PostgresqlConnection.copyOut()"));
131+
return;
132+
}
133+
134+
if (message instanceof CopyInResponse) {
135+
sink.error(new IllegalStateException("COPY ... FROM STDIN statements may only be used with PostgresqlConnection.copyIn()"));
130136
}
131137

132138
} finally {

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@
5757
*/
5858
final class PostgresqlStatement implements io.r2dbc.postgresql.api.PostgresqlStatement {
5959

60-
private static final Predicate<BackendMessage> WINDOW_UNTIL = or(CommandComplete.class::isInstance, EmptyQueryResponse.class::isInstance, ErrorResponse.class::isInstance);
60+
static final Predicate<BackendMessage> WINDOW_UNTIL = or(CommandComplete.class::isInstance, EmptyQueryResponse.class::isInstance, ErrorResponse.class::isInstance);
6161

6262
private final ArrayDeque<Binding> bindings;
6363

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

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,7 @@
1818

1919
import io.netty.buffer.ByteBuf;
2020
import io.r2dbc.postgresql.message.frontend.CancelRequest;
21-
import io.r2dbc.spi.Connection;
22-
import io.r2dbc.spi.IsolationLevel;
23-
import io.r2dbc.spi.R2dbcNonTransientResourceException;
24-
import io.r2dbc.spi.TransactionDefinition;
25-
import io.r2dbc.spi.ValidationDepth;
21+
import io.r2dbc.spi.*;
2622
import org.reactivestreams.Publisher;
2723
import org.reactivestreams.Subscriber;
2824
import reactor.core.publisher.Flux;
@@ -70,6 +66,15 @@ public interface PostgresqlConnection extends Connection {
7066
*/
7167
CopyInBuilder copyIn(String sql);
7268

69+
/**
70+
* Executes one or more PostgreSQL COPY ... TO STDOUT statements and returns the {@link PostgresqlCopyOutResult}s.
71+
* {@link PostgresqlCopyOutResult} objects must be fully consumed to ensure full execution of the {@link Statement}.
72+
*
73+
* @return the {@link PostgresqlCopyOutResult}s, returned by each statement
74+
* @since 1.2
75+
*/
76+
Flux<io.r2dbc.postgresql.api.PostgresqlCopyOutResult> copyOut(String sql);
77+
7378
/**
7479
* Use {@code COPY FROM STDIN} for very fast copying into a database table.
7580
*

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

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
/*
2+
* Copyright 2019 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package io.r2dbc.postgresql.api;
18+
19+
import io.netty.buffer.ByteBuf;
20+
import io.netty.util.ReferenceCounted;
21+
import io.r2dbc.postgresql.message.Format;
22+
import reactor.core.publisher.Flux;
23+
24+
import java.util.Collection;
25+
import java.util.function.BiFunction;
26+
27+
/**
28+
* A result obtained by executing a {@code COPY ... TO STDOUT} statement.
29+
*/
30+
public interface PostgresqlCopyOutResult {
31+
32+
/**
33+
* Returns a mapping of the rows that are the results of a copy-out operation against a database. May be empty if
34+
* the query did not return any rows.
35+
* <p>According to the PostgreSQL documentation, each {@link ByteBuf} will represent a complete row either in
36+
* binary or textual representation, however it is not clear whether that is just an incidental implementation
37+
* quirk or a strict guarantee. In the other direction, there is no such requirement, so robust consumers should
38+
* probably not make any assumptions here.
39+
* <p>The {@link ByteBuf} instances are only valid to use during the {@code mappingFunction} callback, unless they
40+
* are retained (via {@link ReferenceCounted#retain()} and later released. When retaining them, care must be taken
41+
* to ensure that they do eventually get released, as the asynchronous nature of streams can often see messages
42+
* dropped when the stream is canceled downstream.
43+
*
44+
* @param mappingFunction the {@link BiFunction} that maps a {@link ByteBuf} and {@link CopyOutMetadata} to a value
45+
* @param <T> the type of the mapped value
46+
* @return a mapping of the rows that are the results of a copy-out operation against a database
47+
* @throws IllegalArgumentException if {@code mappingFunction} is {@code null}
48+
* @throws IllegalStateException if the result was consumed
49+
*/
50+
<T> Flux<T> map(BiFunction<ByteBuf, CopyOutMetadata, ? extends T> mappingFunction);
51+
52+
/**
53+
* Metadata for a copy-out operation.
54+
*/
55+
interface CopyOutMetadata {
56+
/**
57+
* Get the format of each of the columns.
58+
*
59+
* @return The format of each of the columns
60+
*/
61+
Collection<Format> getColumnFormats();
62+
63+
/**
64+
* Get the overall format of a response.
65+
*
66+
* @return The format.
67+
*/
68+
Format getOverallFormat();
69+
}
70+
}

0 commit comments

Comments
 (0)