Skip to content

Commit f9d0d7f

Browse files
Squirymp911de
Squiry
authored and
mp911de
committed
Handling errors in parse part of extended query flow
[#223]
1 parent 08bc7ea commit f9d0d7f

File tree

7 files changed

+198
-77
lines changed

7 files changed

+198
-77
lines changed

‎src/main/java/io/r2dbc/postgresql/client/Client.java

+23-1
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828

2929
import java.util.Optional;
3030
import java.util.function.Consumer;
31+
import java.util.function.Predicate;
3132

3233
/**
3334
* An abstraction that wraps the networking part of exchanging methods.
@@ -57,7 +58,28 @@ public interface Client {
5758
* @return a {@link Flux} of incoming messages that ends with the end of the frame (i.e. reception of a {@link ReadyForQuery} message.
5859
* @throws IllegalArgumentException if {@code requests} is {@code null}
5960
*/
60-
Flux<BackendMessage> exchange(Publisher<FrontendMessage> requests);
61+
default Flux<BackendMessage> exchange(Publisher<FrontendMessage> requests) {
62+
return this.exchange(it -> it.getClass() == ReadyForQuery.class, requests);
63+
}
64+
65+
/**
66+
* Perform an exchange of messages.
67+
*
68+
* @param requests the publisher of outbound messages
69+
* @return a {@link Flux} of incoming messages that ends with the end of the frame (i.e. reception of a {@link ReadyForQuery} message.
70+
* @throws IllegalArgumentException if {@code requests} is {@code null}
71+
* @since 0.8.1
72+
*/
73+
Flux<BackendMessage> exchange(Predicate<BackendMessage> takeUntil, Publisher<FrontendMessage> requests);
74+
75+
/**
76+
* Send one message without waiting for response.
77+
*
78+
* @param message outbound message
79+
* @throws IllegalArgumentException if {@code message} is {@code null}
80+
* @since 0.8.1
81+
*/
82+
void send(FrontendMessage message);
6183

6284
/**
6385
* Returns the {@link ByteBufAllocator}.

‎src/main/java/io/r2dbc/postgresql/client/ExtendedQueryMessageFlow.java

+37-9
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,16 @@
1919
import io.netty.buffer.Unpooled;
2020
import io.r2dbc.postgresql.message.Format;
2121
import io.r2dbc.postgresql.message.backend.BackendMessage;
22-
import io.r2dbc.postgresql.message.backend.NoData;
23-
import io.r2dbc.postgresql.message.backend.RowDescription;
22+
import io.r2dbc.postgresql.message.backend.CloseComplete;
23+
import io.r2dbc.postgresql.message.backend.ErrorResponse;
24+
import io.r2dbc.postgresql.message.backend.ParseComplete;
25+
import io.r2dbc.postgresql.message.backend.ReadyForQuery;
2426
import io.r2dbc.postgresql.message.frontend.Bind;
2527
import io.r2dbc.postgresql.message.frontend.Close;
2628
import io.r2dbc.postgresql.message.frontend.Describe;
2729
import io.r2dbc.postgresql.message.frontend.Execute;
30+
import io.r2dbc.postgresql.message.frontend.ExecutionType;
31+
import io.r2dbc.postgresql.message.frontend.Flush;
2832
import io.r2dbc.postgresql.message.frontend.FrontendMessage;
2933
import io.r2dbc.postgresql.message.frontend.Parse;
3034
import io.r2dbc.postgresql.message.frontend.Sync;
@@ -36,13 +40,10 @@
3640
import java.util.Collection;
3741
import java.util.Collections;
3842
import java.util.List;
39-
import java.util.function.Predicate;
4043
import java.util.regex.Pattern;
4144

4245
import static io.r2dbc.postgresql.message.frontend.Execute.NO_LIMIT;
4346
import static io.r2dbc.postgresql.message.frontend.ExecutionType.PORTAL;
44-
import static io.r2dbc.postgresql.message.frontend.ExecutionType.STATEMENT;
45-
import static io.r2dbc.postgresql.util.PredicateUtils.or;
4647

4748
/**
4849
* A utility class that encapsulates the <a href="https://www.postgresql.org/docs/current/static/protocol-flow.html#PROTOCOL-FLOW-EXT-QUERY">Extended query</a> message flow.
@@ -54,8 +55,6 @@ public final class ExtendedQueryMessageFlow {
5455
*/
5556
public static final Pattern PARAMETER_SYMBOL = Pattern.compile("\\$([\\d]+)", Pattern.DOTALL);
5657

57-
private static final Predicate<BackendMessage> TAKE_UNTIL = or(RowDescription.class::isInstance, NoData.class::isInstance);
58-
5958
private ExtendedQueryMessageFlow() {
6059
}
6160

@@ -98,8 +97,37 @@ public static Flux<BackendMessage> parse(Client client, String name, String quer
9897
Assert.requireNonNull(query, "query must not be null");
9998
Assert.requireNonNull(types, "types must not be null");
10099

101-
return client.exchange(Flux.just(new Parse(name, types, query), new Describe(name, STATEMENT), Sync.INSTANCE))
102-
.takeUntil(TAKE_UNTIL);
100+
/*
101+
ParseComplete will be received if parse was successful
102+
ReadyForQuery will be received as a response to Sync, which was send in case of error in parsing
103+
*/
104+
return client.exchange(message -> message instanceof ParseComplete || message instanceof ReadyForQuery, Flux.just(new Parse(name, types, query), Flush.INSTANCE))
105+
.doOnNext(message -> {
106+
if (message instanceof ErrorResponse) {
107+
/*
108+
When an error is detected while processing any extended-query message, the backend issues ErrorResponse, then reads and discards messages until a Sync is reached.
109+
So we have to provide Sync message to continue.
110+
*/
111+
client.send(Sync.INSTANCE);
112+
}
113+
});
114+
}
115+
116+
/**
117+
* Execute the close portion of the <a href="https://www.postgresql.org/docs/current/static/protocol-flow.html#PROTOCOL-FLOW-EXT-QUERY">Extended query</a> message flow.
118+
*
119+
* @param client the {@link Client} to exchange messages with
120+
* @param name the name of the statement to close
121+
* @return the messages received in response to this exchange
122+
* @throws IllegalArgumentException if {@code client}, {@code name}, {@code query}, or {@code types} is {@code null}
123+
* @since 0.8.1
124+
*/
125+
public static Flux<BackendMessage> closeStatement(Client client, String name) {
126+
Assert.requireNonNull(client, "client must not be null");
127+
Assert.requireNonNull(name, "name must not be null");
128+
129+
return client.exchange(Flux.just(new Close(name, ExecutionType.STATEMENT), Sync.INSTANCE))
130+
.takeUntil(CloseComplete.class::isInstance);
103131
}
104132

105133
private static Collection<Format> resultFormat(boolean forceBinary) {

‎src/main/java/io/r2dbc/postgresql/client/ReactorNettyClient.java

+67-56
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,6 @@
5252
import reactor.core.publisher.Flux;
5353
import reactor.core.publisher.FluxSink;
5454
import reactor.core.publisher.Mono;
55-
import reactor.core.publisher.MonoSink;
5655
import reactor.core.publisher.SynchronousSink;
5756
import reactor.netty.Connection;
5857
import reactor.netty.resources.ConnectionProvider;
@@ -76,7 +75,7 @@
7675
import java.util.concurrent.atomic.AtomicInteger;
7776
import java.util.concurrent.atomic.AtomicReference;
7877
import java.util.function.Consumer;
79-
import java.util.function.Function;
78+
import java.util.function.Predicate;
8079
import java.util.function.Supplier;
8180

8281
import static io.r2dbc.postgresql.client.TransactionStatus.IDLE;
@@ -104,7 +103,7 @@ public final class ReactorNettyClient implements Client {
104103

105104
private final FluxSink<FrontendMessage> requests = this.requestProcessor.sink();
106105

107-
private final Queue<MonoSink<Flux<BackendMessage>>> responseReceivers = Queues.<MonoSink<Flux<BackendMessage>>>unbounded().get();
106+
private final Queue<ResponseReceiver> responseReceivers = Queues.<ResponseReceiver>unbounded().get();
108107

109108
private final DirectProcessor<NotificationResponse> notificationProcessor = DirectProcessor.create();
110109

@@ -140,21 +139,15 @@ private ReactorNettyClient(Connection connection) {
140139
receiveError.set(throwable);
141140
handleConnectionError(throwable);
142141
})
143-
.windowWhile(it -> it.getClass() != ReadyForQuery.class)
144-
.doOnNext(fluxOfMessages -> {
145-
MonoSink<Flux<BackendMessage>> receiver = this.responseReceivers.poll();
142+
.handle((message, sink) -> {
143+
ResponseReceiver receiver = this.responseReceivers.peek();
146144
if (receiver != null) {
147-
receiver.success(fluxOfMessages.doOnComplete(() -> {
148-
149-
Throwable throwable = receiveError.get();
150-
if (throwable != null) {
151-
throw new PostgresConnectionException(throwable);
152-
}
153-
154-
if (!isConnected()) {
155-
throw EXPECTED.get();
156-
}
157-
}));
145+
if (receiver.takeUntil.test(message)) {
146+
receiver.sink.complete();
147+
this.responseReceivers.poll();
148+
} else {
149+
receiver.sink.next(message);
150+
}
158151
}
159152
})
160153
.doOnComplete(this::handleClose)
@@ -179,6 +172,50 @@ private ReactorNettyClient(Connection connection) {
179172
.subscribe();
180173
}
181174

175+
@Override
176+
public Flux<BackendMessage> exchange(Predicate<BackendMessage> takeUntil, Publisher<FrontendMessage> requests) {
177+
Assert.requireNonNull(requests, "requests must not be null");
178+
179+
return Flux
180+
.create(sink -> {
181+
182+
final AtomicInteger once = new AtomicInteger();
183+
184+
Flux.from(requests)
185+
.subscribe(message -> {
186+
187+
if (!isConnected()) {
188+
ReferenceCountUtil.safeRelease(message);
189+
sink.error(new PostgresConnectionClosedException("Cannot exchange messages because the connection is closed"));
190+
return;
191+
}
192+
193+
if (once.get() == 0 && once.compareAndSet(0, 1)) {
194+
synchronized (this) {
195+
this.responseReceivers.add(new ResponseReceiver(sink, takeUntil));
196+
this.requests.next(message);
197+
}
198+
} else {
199+
this.requests.next(message);
200+
}
201+
202+
}, this.requests::error, () -> {
203+
204+
if (!isConnected()) {
205+
sink.error(new PostgresConnectionClosedException("Cannot exchange messages because the connection is closed"));
206+
}
207+
});
208+
209+
});
210+
}
211+
212+
@Override
213+
public void send(FrontendMessage message) {
214+
Assert.requireNonNull(message, "requests must not be null");
215+
216+
this.requests.next(message);
217+
}
218+
182219
private Mono<Void> resumeError(Throwable throwable) {
183220

184221
handleConnectionError(throwable);
@@ -357,42 +394,12 @@ public Mono<Void> close() {
357394
});
358395
}
359396

360-
@Override
361-
public Flux<BackendMessage> exchange(Publisher<FrontendMessage> requests) {
362-
Assert.requireNonNull(requests, "requests must not be null");
363-
364-
return Mono
365-
.<Flux<BackendMessage>>create(sink -> {
366-
367-
final AtomicInteger once = new AtomicInteger();
368-
369-
Flux.from(requests)
370-
.subscribe(message -> {
371-
372-
if (!isConnected()) {
373-
ReferenceCountUtil.safeRelease(message);
374-
sink.error(new PostgresConnectionClosedException("Cannot exchange messages because the connection is closed"));
375-
return;
376-
}
377-
378-
if (once.get() == 0 && once.compareAndSet(0, 1)) {
379-
synchronized (this) {
380-
this.responseReceivers.add(sink);
381-
this.requests.next(message);
382-
}
383-
} else {
384-
this.requests.next(message);
385-
}
386-
387-
}, this.requests::error, () -> {
388-
389-
if (!isConnected()) {
390-
sink.error(new PostgresConnectionClosedException("Cannot exchange messages because the connection is closed"));
391-
}
392-
});
397+
private void drainError(Supplier<? extends Throwable> supplier) {
398+
ResponseReceiver receiver;
393399

394-
})
395-
.flatMapMany(Function.identity());
400+
while ((receiver = this.responseReceivers.poll()) != null) {
401+
receiver.sink.error(supplier.get());
402+
}
396403
}
397404

398405
@Override
@@ -461,11 +468,15 @@ private void handleConnectionError(Throwable error) {
461468
drainError(() -> new PostgresConnectionException(error));
462469
}
463470

464-
private void drainError(Supplier<? extends Throwable> supplier) {
465-
MonoSink<Flux<BackendMessage>> receiver;
471+
private static class ResponseReceiver {
466472

467-
while ((receiver = this.responseReceivers.poll()) != null) {
468-
receiver.error(supplier.get());
473+
private final FluxSink<BackendMessage> sink;
474+
475+
private final Predicate<BackendMessage> takeUntil;
476+
477+
private ResponseReceiver(FluxSink<BackendMessage> sink, Predicate<BackendMessage> takeUntil) {
478+
this.sink = sink;
479+
this.takeUntil = takeUntil;
469480
}
470481
}
471482

‎src/test/java/io/r2dbc/postgresql/IndefiniteStatementCacheTest.java

+5-7
Original file line numberDiff line numberDiff line change
@@ -22,10 +22,8 @@
2222
import io.r2dbc.postgresql.client.TestClient;
2323
import io.r2dbc.postgresql.message.backend.ErrorResponse;
2424
import io.r2dbc.postgresql.message.backend.ParseComplete;
25-
import io.r2dbc.postgresql.message.frontend.Describe;
26-
import io.r2dbc.postgresql.message.frontend.ExecutionType;
25+
import io.r2dbc.postgresql.message.frontend.Flush;
2726
import io.r2dbc.postgresql.message.frontend.Parse;
28-
import io.r2dbc.postgresql.message.frontend.Sync;
2927
import io.r2dbc.spi.R2dbcNonTransientResourceException;
3028
import org.junit.jupiter.api.Test;
3129
import reactor.core.publisher.Flux;
@@ -50,11 +48,11 @@ void constructorNoClient() {
5048
void getName() {
5149
// @formatter:off
5250
Client client = TestClient.builder()
53-
.expectRequest(new Parse("S_0", Collections.singletonList(100), "test-query"), new Describe("S_0", ExecutionType.STATEMENT), Sync.INSTANCE)
51+
.expectRequest(new Parse("S_0", Collections.singletonList(100), "test-query"), Flush.INSTANCE)
5452
.thenRespond(ParseComplete.INSTANCE)
55-
.expectRequest(new Parse("S_1", Collections.singletonList(200), "test-query"), new Describe("S_1", ExecutionType.STATEMENT), Sync.INSTANCE)
53+
.expectRequest(new Parse("S_1", Collections.singletonList(200), "test-query"), Flush.INSTANCE)
5654
.thenRespond(ParseComplete.INSTANCE)
57-
.expectRequest(new Parse("S_2", Collections.singletonList(200), "test-query-2"), new Describe("S_2", ExecutionType.STATEMENT), Sync.INSTANCE)
55+
.expectRequest(new Parse("S_2", Collections.singletonList(200), "test-query-2"), Flush.INSTANCE)
5856
.thenRespond(ParseComplete.INSTANCE)
5957
.build();
6058
// @formatter:on
@@ -86,7 +84,7 @@ void getName() {
8684
void getNameErrorResponse() {
8785
// @formatter:off
8886
Client client = TestClient.builder()
89-
.expectRequest(new Parse("S_0", Collections.singletonList(100), "test-query"), new Describe("S_0", ExecutionType.STATEMENT), Sync.INSTANCE)
87+
.expectRequest(new Parse("S_0", Collections.singletonList(100), "test-query"), Flush.INSTANCE)
9088
.thenRespond(new ErrorResponse(Collections.emptyList()))
9189
.build();
9290
// @formatter:on

0 commit comments

Comments
 (0)