From 7b2d36caa5ddd18e0ac1e501da1d95ce8a7e5c62 Mon Sep 17 00:00:00 2001 From: doxlik Date: Sat, 12 Sep 2026 16:40:21 +0400 Subject: [PATCH] Support streaming COPY FROM STDIN and COPY TO STDOUT Expose the PostgreSQL COPY subprotocol on PgConnection as raw byte streams, so bulk data transfers no longer have to be materialized in memory. copyOut(String) executes a COPY ... TO STDOUT statement and returns a PgCopyOut, a ReadStream of Buffer whose completion() carries the row count reported by the server. Each CopyData message is emitted as one buffer by default; PgCopyOutOptions.setAggregationThreshold(int) coalesces small messages when throughput matters more than message boundaries. copyIn(String) executes a COPY ... FROM STDIN statement and returns a PgCopyIn, a WriteStream of Buffer honouring writeQueueFull() and drainHandler(). Small writes are batched up to PgCopyInOptions chunkSize, and abort(String) sends CopyFail so the server rolls the copy back. The outcome of a COPY is always the one PostgreSQL reports, never synthesised by the client: completion() resolves at ReadyForQuery, which is also what leaves the connection usable for the next statement. Signed-off-by: doxlik --- vertx-pg-client/src/main/asciidoc/index.adoc | 79 + .../main/java/examples/PgClientExamples.java | 86 + .../java/io/vertx/pgclient/PgConnection.java | 44 + .../main/java/io/vertx/pgclient/PgCopyIn.java | 54 + .../io/vertx/pgclient/PgCopyInOptions.java | 72 + .../java/io/vertx/pgclient/PgCopyOut.java | 70 + .../io/vertx/pgclient/PgCopyOutOptions.java | 73 + .../pgclient/impl/CopyInStreamCommand.java | 34 + .../vertx/pgclient/impl/CopyInStreamImpl.java | 531 +++++ .../pgclient/impl/CopyInStreamInternal.java | 46 + .../io/vertx/pgclient/impl/CopyOutEvent.java | 73 + .../pgclient/impl/CopyOutStreamCommand.java | 42 + .../pgclient/impl/CopyOutStreamImpl.java | 268 +++ .../vertx/pgclient/impl/PgConnectionImpl.java | 201 ++ .../pgclient/impl/PgSocketConnection.java | 5 +- .../pgclient/impl/codec/CopyInHandler.java | 16 + .../codec/CopyInStreamPgCommandMessage.java | 241 +++ .../pgclient/impl/codec/CopyOutHandler.java | 20 + .../codec/CopyOutStreamPgCommandMessage.java | 128 ++ .../io/vertx/pgclient/impl/codec/PgCodec.java | 44 + .../pgclient/impl/codec/PgCommandMessage.java | 6 + .../vertx/pgclient/impl/codec/PgDecoder.java | 74 + .../vertx/pgclient/impl/codec/PgEncoder.java | 36 + .../impl/codec/PgProtocolConstants.java | 5 + .../vertx/tests/pgclient/CopyStreamTest.java | 1075 +++++++++++ .../vertx/tests/pgclient/PgCopyTLSTest.java | 236 +++ .../io/vertx/tests/pgclient/PgCopyTest.java | 1716 +++++++++++++++++ .../tests/pgclient/PgCopyTracingTest.java | 161 ++ .../impl/pool/SqlConnectionPool.java | 6 +- .../sqlclient/impl/tracing/QueryReporter.java | 32 +- .../sqlclient/internal/SqlConnectionBase.java | 6 +- .../spi/protocol/QueryCommandBase.java | 2 +- .../sqlclient/spi/protocol/SqlCommand.java | 26 + 33 files changed, 5489 insertions(+), 19 deletions(-) create mode 100644 vertx-pg-client/src/main/java/io/vertx/pgclient/PgCopyIn.java create mode 100644 vertx-pg-client/src/main/java/io/vertx/pgclient/PgCopyInOptions.java create mode 100644 vertx-pg-client/src/main/java/io/vertx/pgclient/PgCopyOut.java create mode 100644 vertx-pg-client/src/main/java/io/vertx/pgclient/PgCopyOutOptions.java create mode 100644 vertx-pg-client/src/main/java/io/vertx/pgclient/impl/CopyInStreamCommand.java create mode 100644 vertx-pg-client/src/main/java/io/vertx/pgclient/impl/CopyInStreamImpl.java create mode 100644 vertx-pg-client/src/main/java/io/vertx/pgclient/impl/CopyInStreamInternal.java create mode 100644 vertx-pg-client/src/main/java/io/vertx/pgclient/impl/CopyOutEvent.java create mode 100644 vertx-pg-client/src/main/java/io/vertx/pgclient/impl/CopyOutStreamCommand.java create mode 100644 vertx-pg-client/src/main/java/io/vertx/pgclient/impl/CopyOutStreamImpl.java create mode 100644 vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/CopyInHandler.java create mode 100644 vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/CopyInStreamPgCommandMessage.java create mode 100644 vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/CopyOutHandler.java create mode 100644 vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/CopyOutStreamPgCommandMessage.java create mode 100644 vertx-pg-client/src/test/java/io/vertx/tests/pgclient/CopyStreamTest.java create mode 100644 vertx-pg-client/src/test/java/io/vertx/tests/pgclient/PgCopyTLSTest.java create mode 100644 vertx-pg-client/src/test/java/io/vertx/tests/pgclient/PgCopyTest.java create mode 100644 vertx-pg-client/src/test/java/io/vertx/tests/pgclient/PgCopyTracingTest.java create mode 100644 vertx-sql-client/src/main/java/io/vertx/sqlclient/spi/protocol/SqlCommand.java diff --git a/vertx-pg-client/src/main/asciidoc/index.adoc b/vertx-pg-client/src/main/asciidoc/index.adoc index 3c42f9563..6ae36c0d8 100644 --- a/vertx-pg-client/src/main/asciidoc/index.adoc +++ b/vertx-pg-client/src/main/asciidoc/index.adoc @@ -13,6 +13,7 @@ The client is reactive and non-blocking, allowing to handle many database connec * Publish / subscribe using PostgreSQL `NOTIFY/LISTEN` * Batch and cursor * Row streaming +* `COPY` streaming * Command pipelining * RxJava API * Direct memory to object without unnecessary copies @@ -656,6 +657,84 @@ create easily create a string directly from the row set: {@link examples.PgClientExamples#collector02Example} ---- +== Streaming COPY + +PostgreSQL `COPY` transfers bulk data between a table and a stream, without going through +the extended query protocol. The client exposes both directions as raw byte streams, so the +data never has to be materialized in memory. + +Each call accepts exactly one `COPY` statement, and a connection can have a single COPY +operation in progress at a time. + +=== COPY TO STDOUT + +{@link io.vertx.pgclient.PgConnection#copyOut(java.lang.String)} executes a +`COPY ... TO STDOUT` statement and returns a {@link io.vertx.pgclient.PgCopyOut}, a +`ReadStream` of `Buffer`. The future is notified once the server has entered COPY OUT mode. + +[source,$lang] +---- +{@link examples.PgClientExamples#copyOut(io.vertx.pgclient.PgConnection)} +---- + +{@link io.vertx.pgclient.PgCopyOut#completion()} is notified with the number of rows +reported by PostgreSQL. + +The stream is _hot_: like other Vert.x read streams it starts in flowing mode. Consuming it +is what advances the connection, so `completion()` remains pending for as long as the stream +is paused or has no data handler. Call `pause()` before installing a handler when you need +explicit demand control: + +[source,$lang] +---- +{@link examples.PgClientExamples#copyOutPaused(io.vertx.pgclient.PgConnection)} +---- + +By default each PostgreSQL `CopyData` message is emitted as one buffer. When message +boundaries do not matter and throughput does, use +{@link io.vertx.pgclient.PgCopyOutOptions#setAggregationThreshold(int)} to combine small +messages into larger buffers: + +[source,$lang] +---- +{@link examples.PgClientExamples#copyOutWithOptions(io.vertx.pgclient.PgConnection)} +---- + +=== COPY FROM STDIN + +{@link io.vertx.pgclient.PgConnection#copyIn(java.lang.String)} executes a +`COPY ... FROM STDIN` statement and returns a {@link io.vertx.pgclient.PgCopyIn}, a +`WriteStream` of `Buffer`. Write the data, then call `end()` to send `CopyDone`: + +[source,$lang] +---- +{@link examples.PgClientExamples#copyIn(io.vertx.pgclient.PgConnection)} +---- + +The future returned by `end()` is notified when the COPY command itself completes, not +merely when the last buffer has been written. The row count is available from +{@link io.vertx.pgclient.PgCopyIn#completion()}. + +The stream honours `writeQueueFull()` and `drainHandler()` like any other `WriteStream`, so +it can be fed with a `Pipe` or with manual flow control. + +Small writes are combined so that they do not become many small messages on the wire. Use +{@link io.vertx.pgclient.PgCopyInOptions#setChunkSize(int)} to change how much is combined +before a payload is sent; a write larger than that is sent on its own: + +[source,$lang] +---- +{@link examples.PgClientExamples#copyInWithOptions(io.vertx.pgclient.PgConnection)} +---- + +To give up on a copy in progress, call {@link io.vertx.pgclient.PgCopyIn#abort(java.lang.String)}. +It sends a `CopyFail` message and PostgreSQL rolls the copy back: + +[source,$lang] +---- +{@link examples.PgClientExamples#copyInAbort(io.vertx.pgclient.PgConnection)} +---- + == Pub/sub PostgreSQL supports pub/sub communication channels. diff --git a/vertx-pg-client/src/main/java/examples/PgClientExamples.java b/vertx-pg-client/src/main/java/examples/PgClientExamples.java index 9e3069174..778b74a48 100644 --- a/vertx-pg-client/src/main/java/examples/PgClientExamples.java +++ b/vertx-pg-client/src/main/java/examples/PgClientExamples.java @@ -13,6 +13,7 @@ import io.vertx.core.Future; import io.vertx.core.Vertx; +import io.vertx.core.buffer.Buffer; import io.vertx.core.json.JsonObject; import io.vertx.core.net.ClientSSLOptions; import io.vertx.core.net.PemTrustOptions; @@ -824,4 +825,89 @@ public void batchReturning(SqlClient client) { public void pgBouncer(PgConnectOptions connectOptions) { connectOptions.setUseLayer7Proxy(true); } + + public void copyOut(PgConnection connection) { + connection + .copyOut("COPY users TO STDOUT (FORMAT csv)") + .onSuccess(stream -> { + stream.handler(buffer -> { + System.out.println("Received " + buffer.length() + " bytes"); + }); + stream.completion().onSuccess(rows -> { + System.out.println("Copied " + rows + " rows"); + }); + }); + } + + public void copyOutPaused(PgConnection connection) { + connection + .copyOut("COPY users TO STDOUT (FORMAT csv)") + .onSuccess(stream -> { + + // Take control of the demand before any data is emitted + stream.pause(); + + stream.handler(buffer -> { + System.out.println("Received " + buffer.length() + " bytes"); + }); + + // Ask for the next chunk when you are ready for it + stream.fetch(1); + }); + } + + public void copyOutWithOptions(PgConnection connection) { + PgCopyOutOptions options = new PgCopyOutOptions() + .setAggregationThreshold(64 * 1024); + + connection + .copyOut("COPY users TO STDOUT (FORMAT csv)", options) + .onSuccess(stream -> { + stream.handler(buffer -> { + System.out.println("Received " + buffer.length() + " bytes"); + }); + }); + } + + public void copyIn(PgConnection connection) { + connection + .copyIn("COPY users FROM STDIN (FORMAT csv)") + .onSuccess(stream -> { + + stream.write(Buffer.buffer("1,Julien\n")); + stream.write(Buffer.buffer("2,Emad\n")); + + // Sends CopyDone, the future is notified when the COPY command completes + stream.end().onSuccess(v -> { + System.out.println("Copy completed"); + }); + + stream.completion().onSuccess(rows -> { + System.out.println("Copied " + rows + " rows"); + }); + }); + } + + public void copyInWithOptions(PgConnection connection) { + PgCopyInOptions options = new PgCopyInOptions() + .setChunkSize(512 * 1024); + + connection + .copyIn("COPY users FROM STDIN (FORMAT csv)", options) + .onSuccess(stream -> { + stream.write(Buffer.buffer("1,Julien\n")); + stream.end(); + }); + } + + public void copyInAbort(PgConnection connection) { + connection + .copyIn("COPY users FROM STDIN (FORMAT csv)") + .onSuccess(stream -> { + stream.write(Buffer.buffer("1,Julien\n")); + + // Sends CopyFail, the server rolls the COPY back + stream.abort("aborted by the application"); + }); + } } diff --git a/vertx-pg-client/src/main/java/io/vertx/pgclient/PgConnection.java b/vertx-pg-client/src/main/java/io/vertx/pgclient/PgConnection.java index 753148140..07deb2d4a 100644 --- a/vertx-pg-client/src/main/java/io/vertx/pgclient/PgConnection.java +++ b/vertx-pg-client/src/main/java/io/vertx/pgclient/PgConnection.java @@ -135,4 +135,48 @@ static Future connect(Vertx vertx, String connectionUri) { static PgConnection cast(SqlConnection sqlConnection) { return (PgConnection) sqlConnection; } + + /** + * Execute a {@code COPY ... TO STDOUT} statement and stream the raw data produced by the server. + *

+ * The returned future is notified when PostgreSQL has entered COPY OUT mode, with a + * {@link PgCopyOut} emitting the copy data as {@link io.vertx.core.buffer.Buffer} chunks. + * A connection can have a single COPY operation in progress at a time. + * + * @param sql a single {@code COPY ... TO STDOUT} statement + * @return a future notified with the stream or the failure + */ + Future copyOut(String sql); + + /** + * Like {@link #copyOut(String)} but with the given {@code options}. + * + * @param sql a single {@code COPY ... TO STDOUT} statement + * @param options the copy out options + * @return a future notified with the stream or the failure + */ + Future copyOut(String sql, PgCopyOutOptions options); + + /** + * Execute a {@code COPY ... FROM STDIN} statement and stream raw data to the server. + *

+ * The returned future is notified when PostgreSQL has entered COPY IN mode, with a + * {@link PgCopyIn} accepting the copy data as {@link io.vertx.core.buffer.Buffer} chunks. + * Call {@link PgCopyIn#end()} to signal the end of the data, or + * {@link PgCopyIn#abort(String)} to roll the copy back. A connection can have a single + * COPY operation in progress at a time. + * + * @param sql a single {@code COPY ... FROM STDIN} statement + * @return a future notified with the stream or the failure + */ + Future copyIn(String sql); + + /** + * Like {@link #copyIn(String)} but with the given {@code options}. + * + * @param sql a single {@code COPY ... FROM STDIN} statement + * @param options the copy in options + * @return a future notified with the stream or the failure + */ + Future copyIn(String sql, PgCopyInOptions options); } diff --git a/vertx-pg-client/src/main/java/io/vertx/pgclient/PgCopyIn.java b/vertx-pg-client/src/main/java/io/vertx/pgclient/PgCopyIn.java new file mode 100644 index 000000000..240474106 --- /dev/null +++ b/vertx-pg-client/src/main/java/io/vertx/pgclient/PgCopyIn.java @@ -0,0 +1,54 @@ +/* + * Copyright (c) 2011-2026 Contributors to the Eclipse Foundation + * + * This program and the accompanying materials are made available under the + * terms of the Eclipse Public License 2.0 which is available at + * http://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0 + * which is available at https://www.apache.org/licenses/LICENSE-2.0. + * + * SPDX-License-Identifier: EPL-2.0 OR Apache-2.0 + */ + +package io.vertx.pgclient; + +import io.vertx.codegen.annotations.Fluent; +import io.vertx.codegen.annotations.VertxGen; +import io.vertx.core.Future; +import io.vertx.core.Handler; +import io.vertx.core.buffer.Buffer; +import io.vertx.core.streams.WriteStream; + +/** + * COPY FROM STDIN sink as a Vert.x WriteStream. + * + * The stream accepts Buffer chunks and sends them as CopyData messages. + * {@link #completion()} completes with the COPY row count when the server completes the COPY command. + */ +@VertxGen(concrete = false) +public interface PgCopyIn extends WriteStream { + + @Fluent + @Override + PgCopyIn exceptionHandler(Handler handler); + + @Fluent + @Override + PgCopyIn setWriteQueueMaxSize(int maxSize); + + @Fluent + @Override + PgCopyIn drainHandler(Handler handler); + + /** + * Completion of the COPY command. + *

+ * The future succeeds with the row count reported by PostgreSQL, once the server is ready for + * the next command, so the connection can be used again as soon as it is notified. + */ + Future completion(); + + /** + * Fail COPY with a CopyFail message (server will respond with ErrorResponse). + */ + Future abort(String message); +} diff --git a/vertx-pg-client/src/main/java/io/vertx/pgclient/PgCopyInOptions.java b/vertx-pg-client/src/main/java/io/vertx/pgclient/PgCopyInOptions.java new file mode 100644 index 000000000..03586659c --- /dev/null +++ b/vertx-pg-client/src/main/java/io/vertx/pgclient/PgCopyInOptions.java @@ -0,0 +1,72 @@ +/* + * Copyright (c) 2011-2026 Contributors to the Eclipse Foundation + * + * This program and the accompanying materials are made available under the + * terms of the Eclipse Public License 2.0 which is available at + * http://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0 + * which is available at https://www.apache.org/licenses/LICENSE-2.0. + * + * SPDX-License-Identifier: EPL-2.0 OR Apache-2.0 + */ + +package io.vertx.pgclient; + +import io.vertx.codegen.annotations.DataObject; +import io.vertx.codegen.json.annotations.JsonGen; +import io.vertx.core.json.JsonObject; + +/** + * Options for {@link PgConnection#copyIn(String, PgCopyInOptions)}. + */ +@DataObject +@JsonGen(publicConverter = false) +public class PgCopyInOptions { + + /** + * The default target payload size, in bytes = {@code 256 KiB}. + */ + public static final int DEFAULT_CHUNK_SIZE = 256 * 1024; + + private int chunkSize = DEFAULT_CHUNK_SIZE; + + public PgCopyInOptions() { + } + + public PgCopyInOptions(PgCopyInOptions other) { + this.chunkSize = other.chunkSize; + } + + public PgCopyInOptions(JsonObject json) { + PgCopyInOptionsConverter.fromJson(json, this); + } + + /** + * @return the target payload size, in bytes + */ + public int getChunkSize() { + return chunkSize; + } + + /** + * Set the target payload size. + *

+ * Writes are combined until they reach this size before being sent, so that many small writes do + * not become many small messages. A write larger than this size is sent on its own. + * + * @param chunkSize the target payload size, in bytes + * @return a reference to this, so the API can be used fluently + */ + public PgCopyInOptions setChunkSize(int chunkSize) { + if (chunkSize <= 0) { + throw new IllegalArgumentException("chunkSize must be > 0"); + } + this.chunkSize = chunkSize; + return this; + } + + public JsonObject toJson() { + JsonObject json = new JsonObject(); + PgCopyInOptionsConverter.toJson(this, json); + return json; + } +} diff --git a/vertx-pg-client/src/main/java/io/vertx/pgclient/PgCopyOut.java b/vertx-pg-client/src/main/java/io/vertx/pgclient/PgCopyOut.java new file mode 100644 index 000000000..53e19ba37 --- /dev/null +++ b/vertx-pg-client/src/main/java/io/vertx/pgclient/PgCopyOut.java @@ -0,0 +1,70 @@ +/* + * Copyright (c) 2011-2026 Contributors to the Eclipse Foundation + * + * This program and the accompanying materials are made available under the + * terms of the Eclipse Public License 2.0 which is available at + * http://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0 + * which is available at https://www.apache.org/licenses/LICENSE-2.0. + * + * SPDX-License-Identifier: EPL-2.0 OR Apache-2.0 + */ + +package io.vertx.pgclient; + +import io.vertx.codegen.annotations.Fluent; +import io.vertx.codegen.annotations.VertxGen; +import io.vertx.core.Future; +import io.vertx.core.Handler; +import io.vertx.core.buffer.Buffer; +import io.vertx.core.streams.ReadStream; + +/** + * COPY TO STDOUT result as a Vert.x ReadStream. + * + * The stream emits COPY data frames as Buffer chunks. Like other hot Vert.x + * read streams, it is initially in flowing mode. Call {@link #pause()} before + * installing a handler when explicit demand control is required. + * + *

Consuming the stream also advances the PostgreSQL connection. Therefore + * {@link #completion()} may remain pending while the stream is paused or has + * no data handler.

+ * + * {@link #completion()} completes with the COPY row count when the server completes the COPY command. + */ +@VertxGen(concrete = false) +public interface PgCopyOut extends ReadStream { + + @Fluent + @Override + PgCopyOut exceptionHandler(Handler handler); + + @Fluent + @Override + PgCopyOut handler(Handler handler); + + @Fluent + @Override + PgCopyOut pause(); + + @Fluent + @Override + PgCopyOut resume(); + + @Fluent + @Override + PgCopyOut endHandler(Handler endHandler); + + @Fluent + @Override + PgCopyOut fetch(long amount); + + /** + * Completion of the COPY command. + *

+ * The future succeeds with the row count reported by PostgreSQL, once the server is ready for + * the next command, so the connection can be used again as soon as it is notified. When + * transport backpressure is active, the stream must be consumed for the client to observe the + * server's completion. + */ + Future completion(); +} diff --git a/vertx-pg-client/src/main/java/io/vertx/pgclient/PgCopyOutOptions.java b/vertx-pg-client/src/main/java/io/vertx/pgclient/PgCopyOutOptions.java new file mode 100644 index 000000000..b2bf3be68 --- /dev/null +++ b/vertx-pg-client/src/main/java/io/vertx/pgclient/PgCopyOutOptions.java @@ -0,0 +1,73 @@ +/* + * Copyright (c) 2011-2026 Contributors to the Eclipse Foundation + * + * This program and the accompanying materials are made available under the + * terms of the Eclipse Public License 2.0 which is available at + * http://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0 + * which is available at https://www.apache.org/licenses/LICENSE-2.0. + * + * SPDX-License-Identifier: EPL-2.0 OR Apache-2.0 + */ + +package io.vertx.pgclient; + +import io.vertx.codegen.annotations.DataObject; +import io.vertx.codegen.json.annotations.JsonGen; +import io.vertx.core.json.JsonObject; + +/** + * Options for COPY TO STDOUT streams. + */ +@DataObject +@JsonGen(publicConverter = false) +public class PgCopyOutOptions { + + /** + * The default aggregation threshold, in bytes = {@code 1}, each {@code CopyData} message is + * emitted as its own buffer. + */ + public static final int DEFAULT_AGGREGATION_THRESHOLD = 1; + + private int aggregationThreshold = DEFAULT_AGGREGATION_THRESHOLD; + + public PgCopyOutOptions() { + } + + public PgCopyOutOptions(PgCopyOutOptions other) { + this.aggregationThreshold = other.aggregationThreshold; + } + + public PgCopyOutOptions(JsonObject json) { + PgCopyOutOptionsConverter.fromJson(json, this); + } + + /** + * @return the number of accumulated COPY data bytes after which a buffer is emitted + */ + public int getAggregationThreshold() { + return aggregationThreshold; + } + + /** + * Set the number of accumulated COPY data bytes after which a buffer is emitted. + *

+ * A value of {@code 1} disables additional aggregation and emits each PostgreSQL + * {@code CopyData} message as a buffer. COPY data buffers are not row-aligned. + * + * @param aggregationThreshold the aggregation threshold, in bytes + * @return this options instance + */ + public PgCopyOutOptions setAggregationThreshold(int aggregationThreshold) { + if (aggregationThreshold <= 0) { + throw new IllegalArgumentException("aggregationThreshold must be > 0"); + } + this.aggregationThreshold = aggregationThreshold; + return this; + } + + public JsonObject toJson() { + JsonObject json = new JsonObject(); + PgCopyOutOptionsConverter.toJson(this, json); + return json; + } +} diff --git a/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/CopyInStreamCommand.java b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/CopyInStreamCommand.java new file mode 100644 index 000000000..fd81c56a9 --- /dev/null +++ b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/CopyInStreamCommand.java @@ -0,0 +1,34 @@ +/* + * Copyright (c) 2011-2026 Contributors to the Eclipse Foundation + * + * This program and the accompanying materials are made available under the + * terms of the Eclipse Public License 2.0 which is available at + * http://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0 + * which is available at https://www.apache.org/licenses/LICENSE-2.0. + * + * SPDX-License-Identifier: EPL-2.0 OR Apache-2.0 + */ + +package io.vertx.pgclient.impl; + +import io.vertx.sqlclient.spi.protocol.CommandBase; +import io.vertx.sqlclient.spi.protocol.SqlCommand; + +public final class CopyInStreamCommand extends CommandBase implements SqlCommand { + private final String sql; + private final CopyInStreamInternal in; + + public CopyInStreamCommand(String sql, CopyInStreamInternal in) { + this.sql = sql; + this.in = in; + } + + @Override + public String sql() { + return sql; + } + + public CopyInStreamInternal in() { + return in; + } +} diff --git a/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/CopyInStreamImpl.java b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/CopyInStreamImpl.java new file mode 100644 index 000000000..44cd65daa --- /dev/null +++ b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/CopyInStreamImpl.java @@ -0,0 +1,531 @@ +/* + * Copyright (c) 2011-2026 Contributors to the Eclipse Foundation + * + * This program and the accompanying materials are made available under the + * terms of the Eclipse Public License 2.0 which is available at + * http://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0 + * which is available at https://www.apache.org/licenses/LICENSE-2.0. + * + * SPDX-License-Identifier: EPL-2.0 OR Apache-2.0 + */ + +package io.vertx.pgclient.impl; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.CompositeByteBuf; +import io.netty.buffer.Unpooled; +import io.vertx.codegen.annotations.Nullable; +import io.vertx.core.Future; +import io.vertx.core.Handler; +import io.vertx.core.Promise; +import io.vertx.core.buffer.Buffer; +import io.vertx.core.internal.ContextInternal; +import io.vertx.core.internal.buffer.BufferInternal; +import io.vertx.core.streams.WriteStream; +import io.vertx.pgclient.PgCopyIn; +import io.vertx.pgclient.PgCopyInOptions; + +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Deque; +import java.util.List; + +public final class CopyInStreamImpl implements CopyInStreamInternal, WriteStream { + + private static final class PendingFrame { + final ByteBuf buf; + final int bytes; + final List> writePromises; + + PendingFrame(ByteBuf buf, int bytes, List> writePromises) { + this.buf = buf; + this.bytes = bytes; + this.writePromises = writePromises; + } + + void completeWrites() { + if (writePromises == null) return; + for (Promise p : writePromises) { + p.tryComplete(); + } + } + + void failWrites(Throwable t) { + if (writePromises == null) return; + for (Promise p : writePromises) { + p.tryFail(t); + } + } + } + + /** + * Components the accumulating buffer may hold before Netty consolidates it. + *

+ * {@link Unpooled#compositeBuffer()} allows only 16, and every add beyond that copies the whole + * payload accumulated so far, which turns coalescing many small writes into quadratic copying. + * The bound is kept well below the number of buffers a gathering write can carry so that the + * frame still leaves as a single vectored write. + */ + private static final int MAX_CHUNK_COMPONENTS = 1024; + + private final ContextInternal context; + + private final Promise completion; + + private final Promise ready; + + private final int chunkBytes; + + /** Frames built and waiting for the transport to accept them. */ + private final Deque pending = new ArrayDeque<>(); + + private int pendingBytes; + + /** The frame currently being accumulated from application writes. */ + private CompositeByteBuf building; + + private int buildingBytes; + + private List> buildingWritePromises; + + private boolean draining; + + private boolean drainScheduled; + + private Sink sink; + + private int writeQueueMaxSize = 8 * 1024 * 1024; + + private Handler drainHandler; + + private Handler exceptionHandler; + + private boolean wasFull; + + private boolean ended; + + private Throwable failure; + + private boolean copyDoneSent; + + private boolean copyFailSent; + + private String pendingCopyFailMessage; + + public CopyInStreamImpl(ContextInternal context, PgCopyInOptions options) { + this.context = context; + this.completion = context.promise(); + this.ready = context.promise(); + this.chunkBytes = options.getChunkSize(); + } + + @Override + public Future write(Buffer data) { + // Not context.promise(): completed on the connection's event loop anyway, and a context bound + // future costs an isRunningOnContext check plus beginDispatch/endDispatch on every write + Promise promise = Promise.promise(); + + context.emit(promise, p -> { + if (failure != null) { + p.fail(failure); + return; + } + + if (ended) { + p.fail(new IllegalStateException("COPY IN already ended")); + return; + } + + if (data == null) { + p.fail(new NullPointerException("data")); + return; + } + + ByteBuf part = retainByteBuf(data); + + int sz = part.readableBytes(); + if (sz == 0) { + part.release(); + p.complete(); + return; + } + + ensureBuilding(); + + building.addComponent(true, part); + + buildingBytes += sz; + + buildingWritePromises.add(p); + + if (buildingBytes >= chunkBytes || writeQueueFull()) { + enqueueBuilt(); + } + + scheduleDrain(); + }); + + return promise.future(); + } + + @Override + public boolean writeQueueFull() { + return (pendingBytes + buildingBytes) >= writeQueueMaxSize; + } + + @Override + public Future end() { + Promise ignored = context.promise(); + + context.emit(ignored, p -> { + if (ended) { + p.complete(); + return; + } + + ended = true; + + if (buildingBytes > 0) { + enqueueBuilt(); + } + + scheduleDrain(); + p.complete(); + }); + + return completion.future().mapEmpty(); + } + + @Override + public Future completion() { + return completion.future(); + } + + @Override + public Future readyFuture() { + return ready.future(); + } + + @Override + public CopyInStreamImpl exceptionHandler(@Nullable Handler handler) { + this.exceptionHandler = handler; + return this; + } + + @Override + public CopyInStreamImpl setWriteQueueMaxSize(int maxSize) { + if (maxSize < 1) { + maxSize = 1; + } + this.writeQueueMaxSize = maxSize; + + Sink s = this.sink; + if (s != null) { + s.setWatermarks(maxSize); + } + + return this; + } + + @Override + public CopyInStreamImpl drainHandler(@Nullable Handler handler) { + this.drainHandler = handler; + return this; + } + + @Override + public Future abort(String message) { + Promise promise = context.promise(); + + context.emit(promise, p -> { + if (failure != null) { + p.fail(failure); + return; + } + + if (ended) { + p.fail(new IllegalStateException("COPY IN already ended")); + return; + } + + String msg = message != null ? message : "COPY IN aborted"; + pendingCopyFailMessage = msg; + + // Stop the stream and send CopyFail, but leave completion() to the server: it reports the + // real cause when the COPY had already failed for a reason of its own. + failLocally(new IllegalStateException(msg), true); + + p.complete(); + }); + + return promise.future(); + } + + @Override + public void attachSink(Sink sink) { + context.runOnContext(v -> { + + Sink prev = this.sink; + if (prev != null) { + prev.detach(); + } + + if (failure != null) { + if (!copyFailSent) { + copyFailSent = true; + try { + sink.writeCopyFail(pendingCopyFailMessage != null ? pendingCopyFailMessage : safeMsg(failure)); + } catch (Throwable ignore) { + } + } + sink.detach(); + this.sink = null; + return; + } + + this.sink = sink; + + sink.setWatermarks(writeQueueMaxSize); + + sink.onWritable(this::scheduleDrain); + + ready.tryComplete(this); + + syncFullState(); + }); + } + + @Override + public void detachSinkIfAny() { + context.runOnContext(v -> { + Sink s = this.sink; + this.sink = null; + + if (s != null) { + s.detach(); + } + }); + } + + @Override + public void completeFromServer(int rowCount) { + context.runOnContext(v -> completion.tryComplete(rowCount)); + } + + @Override + public void failFromServer(Throwable t) { + context.runOnContext(v -> { + failLocally(t, false); + // Settles completion() even when a local abort already stopped the stream + completion.tryFail(t); + }); + } + + private void notifyIfNoLongerFull() { + boolean fullNow = writeQueueFull(); + + if (wasFull && !fullNow) { + Handler dh = drainHandler; + if (dh != null) { + + context.runOnContext(v -> dh.handle(null)); + } + } + + wasFull = fullNow; + } + + private void syncFullState() { + wasFull = writeQueueFull(); + } + + private void failNow(Throwable t, boolean sendCopyFail) { + failLocally(t, sendCopyFail); + completion.tryFail(t); + } + + /** + * Move the stream to its failed state and notify the application, without deciding the outcome of + * the COPY command itself. + */ + private void failLocally(Throwable t, boolean sendCopyFail) { + if (failure != null) { + return; + } + + failure = t; + + ready.tryFail(t); + + if (sendCopyFail) { + Sink s = sink; + if (s != null && !copyFailSent) { + copyFailSent = true; + try { + s.writeCopyFail(pendingCopyFailMessage != null ? pendingCopyFailMessage : safeMsg(t)); + } catch (Throwable ignore) { + + } + } + } + + failAllPendingPromises(t); + releaseAllPending(); + + Handler eh = exceptionHandler; + if (eh != null) { + context.runOnContext(v -> eh.handle(t)); + } + } + + private static ByteBuf retainByteBuf(Buffer data) { + ByteBuf bb = ((BufferInternal) data).getByteBuf(); + return bb.retainedDuplicate(); + } + + private static String safeMsg(Throwable t) { + String m = t.getMessage(); + return m != null ? m : t.getClass().getName(); + } + + private void scheduleDrain() { + if (drainScheduled) return; + drainScheduled = true; + + context.runOnContext(v -> { + drainScheduled = false; + // Batch writes made in the same turn, but never strand a partial frame. + if (failure == null && buildingBytes > 0) { + enqueueBuilt(); + } + doDrain(); + }); + } + + private void doDrain() { + if (draining) return; + + draining = true; + try { + Sink s = sink; + + if (s == null) { + return; + } + + if (failure != null) { + return; + } + + while (!pending.isEmpty() && s.isWritable()) { + PendingFrame frame = pending.pollFirst(); + pendingBytes -= frame.bytes; + + writeFrame(s, frame); + + if (failure != null) { + return; + } + } + + notifyIfNoLongerFull(); + + if (ended && !copyDoneSent && pending.isEmpty()) { + copyDoneSent = true; + s.writeCopyDone(); + } + + } catch (Throwable t) { + failNow(t, true); + } finally { + draining = false; + } + } + + /** + * Hand the frame to the transport and settle the writes it carries. + *

+ * The write futures resolve once the frame is accepted by the transport, which only happens + * while the channel is writable, so back pressure is preserved. Whether the COPY itself + * succeeded is reported by {@link #completion()}, the only answer that is meaningful for a + * statement PostgreSQL applies atomically. + */ + private void writeFrame(Sink s, PendingFrame frame) { + try { + s.writeCopyData(frame.buf); + + frame.completeWrites(); + } catch (Throwable t) { + + try { + frame.buf.release(); + } catch (Throwable ignore) { + } + + frame.failWrites(t); + failNow(t, true); + } + } + + private void ensureBuilding() { + if (building == null) { + building = Unpooled.compositeBuffer(MAX_CHUNK_COMPONENTS); + buildingBytes = 0; + buildingWritePromises = new ArrayList<>(8); + } + } + + private void enqueueBuilt() { + if (building == null || buildingBytes == 0) { + return; + } + + ByteBuf payload = building; + int sz = buildingBytes; + List> writePromises = buildingWritePromises; + + building = null; + buildingBytes = 0; + buildingWritePromises = null; + + pending.addLast(new PendingFrame(payload, sz, writePromises)); + pendingBytes += sz; + + wasFull = writeQueueFull(); + } + + private void failAllPendingPromises(Throwable t) { + if (buildingWritePromises != null) { + for (Promise p : buildingWritePromises) { + p.tryFail(t); + } + buildingWritePromises = null; + } + + for (PendingFrame f : pending) { + f.failWrites(t); + } + } + + private void releaseAllPending() { + if (building != null) { + try { + building.release(); + } catch (Throwable ignore) { + } + building = null; + buildingBytes = 0; + buildingWritePromises = null; + } + + while (!pending.isEmpty()) { + PendingFrame f = pending.pollFirst(); + try { + f.buf.release(); + } catch (Throwable ignore) { + } + } + + pendingBytes = 0; + } +} diff --git a/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/CopyInStreamInternal.java b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/CopyInStreamInternal.java new file mode 100644 index 000000000..66b2398d8 --- /dev/null +++ b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/CopyInStreamInternal.java @@ -0,0 +1,46 @@ +/* + * Copyright (c) 2011-2026 Contributors to the Eclipse Foundation + * + * This program and the accompanying materials are made available under the + * terms of the Eclipse Public License 2.0 which is available at + * http://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0 + * which is available at https://www.apache.org/licenses/LICENSE-2.0. + * + * SPDX-License-Identifier: EPL-2.0 OR Apache-2.0 + */ + +package io.vertx.pgclient.impl; + +import io.netty.buffer.ByteBuf; +import io.vertx.core.Future; +import io.vertx.pgclient.PgCopyIn; + +public interface CopyInStreamInternal extends PgCopyIn { + + interface Sink { + + boolean isWritable(); + + void onWritable(Runnable cb); + + void setWatermarks(int maxBytes); + + void writeCopyData(ByteBuf buf); + + void writeCopyDone(); + + void writeCopyFail(String message); + + void detach(); + } + + void attachSink(Sink sink); + + void detachSinkIfAny(); + + void completeFromServer(int rowCount); + + void failFromServer(Throwable t); + + Future readyFuture(); +} diff --git a/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/CopyOutEvent.java b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/CopyOutEvent.java new file mode 100644 index 000000000..b46bfeb51 --- /dev/null +++ b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/CopyOutEvent.java @@ -0,0 +1,73 @@ +/* + * Copyright (c) 2011-2026 Contributors to the Eclipse Foundation + * + * This program and the accompanying materials are made available under the + * terms of the Eclipse Public License 2.0 which is available at + * http://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0 + * which is available at https://www.apache.org/licenses/LICENSE-2.0. + * + * SPDX-License-Identifier: EPL-2.0 OR Apache-2.0 + */ + +package io.vertx.pgclient.impl; + +import io.vertx.core.buffer.Buffer; + +import java.util.Objects; + +/** + * An inbound COPY OUT event. + * + *

The PostgreSQL decoder emits these events through the regular Vert.x + * connection pipeline so the socket's existing read queue remains the sole + * owner of Netty auto-read state.

+ */ +public final class CopyOutEvent { + + private enum Kind { + /** A CopyData message. */ + DATA, + /** CommandComplete, no more data will follow. */ + END, + /** ReadyForQuery, the COPY command is done and the connection is free. */ + COMPLETED + } + + private final CopyOutStreamImpl stream; + private final Kind kind; + private final Buffer data; + private final int rowCount; + + private CopyOutEvent(CopyOutStreamImpl stream, Kind kind, Buffer data, int rowCount) { + this.stream = Objects.requireNonNull(stream, "stream"); + this.kind = kind; + this.data = data; + this.rowCount = rowCount; + } + + public static CopyOutEvent data(CopyOutStreamImpl stream, Buffer data) { + return new CopyOutEvent(stream, Kind.DATA, Objects.requireNonNull(data, "data"), -1); + } + + public static CopyOutEvent end(CopyOutStreamImpl stream, int rowCount) { + return new CopyOutEvent(stream, Kind.END, null, rowCount); + } + + public static CopyOutEvent completed(CopyOutStreamImpl stream) { + return new CopyOutEvent(stream, Kind.COMPLETED, null, -1); + } + + void dispatch() { + switch (kind) { + case DATA: + stream.emit(data); + break; + case END: + stream.end(rowCount); + break; + case COMPLETED: + stream.commandCompleted(); + break; + } + } +} diff --git a/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/CopyOutStreamCommand.java b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/CopyOutStreamCommand.java new file mode 100644 index 000000000..26f047284 --- /dev/null +++ b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/CopyOutStreamCommand.java @@ -0,0 +1,42 @@ +/* + * Copyright (c) 2011-2026 Contributors to the Eclipse Foundation + * + * This program and the accompanying materials are made available under the + * terms of the Eclipse Public License 2.0 which is available at + * http://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0 + * which is available at https://www.apache.org/licenses/LICENSE-2.0. + * + * SPDX-License-Identifier: EPL-2.0 OR Apache-2.0 + */ + +package io.vertx.pgclient.impl; + +import io.vertx.pgclient.PgCopyOutOptions; +import io.vertx.sqlclient.spi.protocol.CommandBase; +import io.vertx.sqlclient.spi.protocol.SqlCommand; + +public final class CopyOutStreamCommand extends CommandBase implements SqlCommand { + + private final String sql; + private final CopyOutStreamImpl out; + private final PgCopyOutOptions options; + + public CopyOutStreamCommand(String sql, CopyOutStreamImpl out, PgCopyOutOptions options) { + this.sql = sql; + this.out = out; + this.options = options; + } + + @Override + public String sql() { + return sql; + } + + public CopyOutStreamImpl out() { + return out; + } + + public PgCopyOutOptions options() { + return options; + } +} diff --git a/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/CopyOutStreamImpl.java b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/CopyOutStreamImpl.java new file mode 100644 index 000000000..d8d245394 --- /dev/null +++ b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/CopyOutStreamImpl.java @@ -0,0 +1,268 @@ +/* + * Copyright (c) 2011-2026 Contributors to the Eclipse Foundation + * + * This program and the accompanying materials are made available under the + * terms of the Eclipse Public License 2.0 which is available at + * http://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0 + * which is available at https://www.apache.org/licenses/LICENSE-2.0. + * + * SPDX-License-Identifier: EPL-2.0 OR Apache-2.0 + */ + +package io.vertx.pgclient.impl; + +import io.vertx.core.internal.ContextInternal; +import io.vertx.core.Future; +import io.vertx.core.Handler; +import io.vertx.core.Promise; +import io.vertx.core.buffer.Buffer; +import io.vertx.core.streams.ReadStream; +import io.vertx.pgclient.PgCopyOut; + +import java.util.ArrayDeque; +import java.util.Objects; +import java.util.Queue; + +public final class CopyOutStreamImpl implements PgCopyOut { + + private final ContextInternal ctx; + private final ReadStream upstream; + private final Promise completion; + private final Promise ready; + + private Handler exceptionHandler; + private Handler handler; + private Handler endHandler; + + private boolean paused; + private int rowCount; + private long demand = Long.MAX_VALUE; + private boolean ended; + private boolean endReached; + private boolean endHandlerNotified; + private boolean discarding; + private boolean draining; + private boolean upstreamPaused; + private Throwable failure; + + // Buffers events already dispatched when the upstream stream was paused. + // The upstream Vert.x read queue owns the actual backpressure watermarks. + private final Queue queue = new ArrayDeque<>(); + + public CopyOutStreamImpl(ContextInternal ctx, ReadStream upstream) { + this.ctx = Objects.requireNonNull(ctx, "ctx"); + this.upstream = Objects.requireNonNull(upstream, "upstream"); + this.completion = ctx.promise(); + this.ready = ctx.promise(); + } + + public void readyFromServer() { + // Called from the codec on the connection's event loop, so settle it now: deferring would let + // ReadyForQuery be processed first and make the stream look like it never started. + ready.tryComplete(this); + } + + public Future readyFuture() { + return ready.future(); + } + + public void emit(Buffer buf) { + if (ended) return; + if (buf == null) return; + + if (handler == null || paused || demand == 0) { + queue.add(buf); + updateUpstream(); + return; + } + + deliver(buf); + updateUpstream(); + } + + /** + * The server sent CommandComplete, no more data will be emitted. The COPY row count is held + * until the server is ready for the next command, see {@link #commandCompleted()}. + */ + public void end(int rowCount) { + if (ended) return; + ended = true; + this.rowCount = rowCount; + drain(); + } + + /** + * The server sent ReadyForQuery, the COPY command is done and the connection is free again. + */ + public void commandCompleted() { + completion.tryComplete(rowCount); + } + + public void discard(Throwable t) { + discarding = true; + queue.clear(); + resumeUpstream(); + fail(t); + } + + public boolean isDiscarding() { + return discarding; + } + + public void fail(Throwable t) { + // Always settle the futures the caller holds. A statement that completed without entering copy + // mode ends the stream without ever handing one out, and the data can be drained before the + // server acknowledges the command, leaving the completion for a ReadyForQuery that never comes. + ready.tryFail(t); + completion.tryFail(t); + if (failure != null || endReached) return; + failure = Objects.requireNonNull(t, "t"); + discarding = true; + ended = true; + queue.clear(); + resumeUpstream(); + Handler h = exceptionHandler; + if (h != null) { + ctx.runOnContext(v -> h.handle(t)); + } + } + + @Override + public CopyOutStreamImpl exceptionHandler(Handler handler) { + this.exceptionHandler = handler; + if (failure != null && handler != null) { + ctx.runOnContext(v -> handler.handle(failure)); + } + return this; + } + + @Override + public CopyOutStreamImpl handler(Handler handler) { + this.handler = handler; + drain(); + return this; + } + + @Override + public CopyOutStreamImpl pause() { + paused = true; + demand = 0; + updateUpstream(); + return this; + } + + @Override + public CopyOutStreamImpl resume() { + paused = false; + demand = Long.MAX_VALUE; + drain(); + return this; + } + + @Override + public CopyOutStreamImpl fetch(long amount) { + if (amount < 0) throw new IllegalArgumentException("amount < 0"); + paused = false; + + if (demand != Long.MAX_VALUE) { + long next = demand + amount; + if (next < 0) { + demand = Long.MAX_VALUE; + } else { + demand = next; + } + } + + drain(); + return this; + } + + @Override + public CopyOutStreamImpl endHandler(Handler endHandler) { + this.endHandler = endHandler; + notifyEndHandler(); + return this; + } + + @Override + public Future completion() { + return completion.future(); + } + + private void drain() { + if (draining) { + return; + } + if (handler == null || paused || demand == 0) { + updateUpstream(); + signalEndIfDrained(); + return; + } + draining = true; + try { + while (!queue.isEmpty() && handler != null && !paused && demand != 0) { + Buffer buf = queue.poll(); + deliver(buf); + } + } finally { + draining = false; + updateUpstream(); + signalEndIfDrained(); + } + } + + private void deliver(Buffer buf) { + Handler h = handler; + if (h == null) { + queue.add(buf); + return; + } + + if (demand != Long.MAX_VALUE) { + demand--; + } + + try { + h.handle(buf); + } catch (Throwable t) { + fail(t); + } + } + + private void updateUpstream() { + if (ended || (handler != null && !paused && demand != 0)) { + resumeUpstream(); + } else { + pauseUpstream(); + } + } + + private void pauseUpstream() { + if (!upstreamPaused) { + upstreamPaused = true; + upstream.pause(); + } + } + + private void resumeUpstream() { + if (upstreamPaused) { + upstreamPaused = false; + upstream.resume(); + } + } + + private void signalEndIfDrained() { + if (ended && queue.isEmpty() && failure == null) { + endReached = true; + notifyEndHandler(); + } + } + + private void notifyEndHandler() { + Handler h = endHandler; + if (endReached && !endHandlerNotified && h != null) { + endHandlerNotified = true; + ctx.runOnContext(v -> h.handle(null)); + } + } +} diff --git a/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/PgConnectionImpl.java b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/PgConnectionImpl.java index d81dbdd02..7ac47b8ba 100644 --- a/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/PgConnectionImpl.java +++ b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/PgConnectionImpl.java @@ -20,6 +20,10 @@ import io.vertx.core.internal.ContextInternal; import io.vertx.pgclient.PgConnectOptions; import io.vertx.pgclient.PgConnection; +import io.vertx.pgclient.PgCopyIn; +import io.vertx.pgclient.PgCopyInOptions; +import io.vertx.pgclient.PgCopyOut; +import io.vertx.pgclient.PgCopyOutOptions; import io.vertx.pgclient.PgNotice; import io.vertx.pgclient.PgNotification; import io.vertx.pgclient.impl.codec.NoticeResponse; @@ -29,6 +33,9 @@ import io.vertx.sqlclient.internal.SqlConnectionBase; import io.vertx.sqlclient.spi.connection.Connection; +import java.util.Locale; +import java.util.function.Supplier; + public class PgConnectionImpl extends SqlConnectionBase implements PgConnection { public static Future connect(ContextInternal context, PgConnectOptions options) { @@ -48,6 +55,40 @@ public static Future connect(ContextInternal context, PgConnectOpt private volatile Handler notificationHandler; private volatile Handler noticeHandler; + /** + * The COPY operation currently owning the connection, {@code null} when there is none. + *

+ * Like {@link #tx} this is confined to the connection context and cleared by the operation itself + * when the server is done with it, a connection carries at most one of them at a time. + */ + private ActiveCopy activeCopy; + + /** + * A COPY operation in progress, remembered so that closing the connection leaves copy mode + * instead of recycling a connection PostgreSQL still considers to be copying. + */ + private static final class ActiveCopy { + + private final Supplier> leaveCopyMode; + private final Future commandCompletion; + + ActiveCopy(Supplier> leaveCopyMode, Future commandCompletion) { + this.leaveCopyMode = leaveCopyMode; + this.commandCompletion = commandCompletion; + } + + /** + * Tell the server to end the copy and wait for it to acknowledge, so the connection is usable + * again. The command is expected to fail, aborting a COPY is reported as an error by the + * server, and the connection is closing either way. + */ + Future cancel() { + return leaveCopyMode.get() + .transform(ignored -> commandCompletion) + .otherwiseEmpty(); + } + } + public PgConnectionImpl(PgConnectionFactory factory, ContextInternal context, Connection conn) { super(context, factory, conn, PgDriver.INSTANCE); } @@ -56,6 +97,83 @@ public PgConnectionImpl(PgConnectionFactory factory, ContextInternal context, Co super(context, factory, conn, PgDriver.INSTANCE, registerCleanup); } + /** + * {@inheritDoc} + *

+ * A COPY in progress is left before the connection goes, so that a pooled connection is never + * recycled while PostgreSQL still considers it to be copying. COPY IN is given up with + * {@code CopyFail}, COPY OUT has no such message, the remaining rows are read and dropped, so + * closing during a large COPY OUT waits for the server to finish sending it. + */ + @Override + public Future close() { + ActiveCopy copy = activeCopy; + if (copy == null) { + return super.close(); + } + return copy.cancel().compose(v -> PgConnectionImpl.super.close()); + } + + /** + * Remember {@code copy} until the server is done with it, so that the connection is not handed + * to another COPY, or back to a pool, while it is still in copy mode. + */ + private void beginCopy(ActiveCopy copy) { + activeCopy = copy; + copy.commandCompletion.onComplete(ar -> { + if (activeCopy == copy) { + activeCopy = null; + } + }); + } + + /** + * COPY takes over the connection, so the statement has to be the right kind of COPY, and it has + * to be the only one: a trailing statement would run after the transfer and its + * {@code CommandComplete} would be mistaken for the COPY's, reporting a nonsense row count. + * + * @return a failed future to return to the caller, or {@code null} when the statement is fine + */ + private Future checkCopySql(String sql, String direction) { + if (sql == null) { + return context.failedFuture(new NullPointerException("sql")); + } + String trimmed = sql.trim(); + // Keywords are separated by any run of whitespace, a statement may well be wrapped over lines + String normalized = trimmed.toUpperCase(Locale.ROOT).replaceAll("\\s+", " "); + if (!normalized.startsWith("COPY") || !normalized.contains(direction)) { + return context.failedFuture(new IllegalArgumentException("Not a COPY " + direction + " statement: " + sql)); + } + if (hasTrailingStatement(trimmed)) { + return context.failedFuture(new IllegalArgumentException("COPY accepts a single statement: " + sql)); + } + return null; + } + + /** + * Whether something follows the first statement separator, {@code COPY t FROM STDIN (DELIMITER ';')} + * is a single statement, the semicolon belongs to the option rather than separating statements. + */ + private static boolean hasTrailingStatement(String sql) { + boolean inString = false; + boolean inIdentifier = false; + for (int i = 0; i < sql.length(); i++) { + char c = sql.charAt(i); + if (c == '\'' && !inIdentifier) { + inString = !inString; + } else if (c == '"' && !inString) { + inIdentifier = !inIdentifier; + } else if (c == ';' && !inString && !inIdentifier) { + return !sql.substring(i + 1).trim().isEmpty(); + } + } + return false; + } + + private Future copyAlreadyInProgress() { + return context.failedFuture(new IllegalStateException("A COPY operation is already in progress on this connection")); + } + @Override public PgConnection notificationHandler(Handler handler) { notificationHandler = handler; @@ -133,4 +251,87 @@ public Future cancelRequest() { }); return promise.future(); } + + @Override + public Future copyOut(String sql) { + return copyOut(sql, null); + } + + @Override + public Future copyOut(String sql, PgCopyOutOptions options) { + Future rejected = checkCopySql(sql, "TO STDOUT"); + if (rejected != null) { + return rejected; + } + + if (activeCopy != null) { + return copyAlreadyInProgress(); + } + + PgCopyOutOptions finalOptions = options != null ? options : new PgCopyOutOptions(); + Promise completion = context.promise(); + Future commandCompletion = completion.future(); + PgSocketConnection actual = (PgSocketConnection) conn.unwrap(); + CopyOutStreamImpl out = new CopyOutStreamImpl(context, actual.socket()); + beginCopy(new ActiveCopy(() -> { + out.discard(new IllegalStateException("Connection closed")); + return Future.succeededFuture(); + }, commandCompletion)); + + try { + schedule(new CopyOutStreamCommand(sql, out, finalOptions), completion); + } catch (RuntimeException e) { + completion.tryFail(e); + } + commandCompletion.onFailure(out::fail); + // A COPY that ran without entering copy mode leaves nothing to hand out, so say so rather than + // leaving the caller holding a future that never resolves. + commandCompletion.onSuccess(v -> { + if (!out.readyFuture().isComplete()) { + out.fail(new IllegalStateException("COPY did not start: " + sql)); + } + }); + + return out.readyFuture(); + } + + @Override + public Future copyIn(String sql) { + return copyIn(sql, null); + } + + @Override + public Future copyIn(String sql, PgCopyInOptions options) { + Future rejected = checkCopySql(sql, "FROM STDIN"); + if (rejected != null) { + return rejected; + } + + if (activeCopy != null) { + return copyAlreadyInProgress(); + } + + PgCopyInOptions finalOptions = options != null ? options : new PgCopyInOptions(); + Promise completion = context.promise(); + Future commandCompletion = completion.future(); + + CopyInStreamInternal in = new CopyInStreamImpl(context, finalOptions); + beginCopy(new ActiveCopy(() -> in.abort("Connection closed"), commandCompletion)); + + try { + schedule(new CopyInStreamCommand(sql, in), completion); + } catch (RuntimeException e) { + completion.tryFail(e); + } + commandCompletion.onFailure(in::failFromServer); + // A COPY that ran without entering copy mode leaves nothing to hand out, so say so rather than + // leaving the caller holding a future that never resolves. + commandCompletion.onSuccess(v -> { + if (!in.readyFuture().isComplete()) { + in.failFromServer(new IllegalStateException("COPY did not start: " + sql)); + } + }); + + return in.readyFuture(); + } } diff --git a/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/PgSocketConnection.java b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/PgSocketConnection.java index 24f9e3edf..70d0b1057 100644 --- a/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/PgSocketConnection.java +++ b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/PgSocketConnection.java @@ -82,6 +82,7 @@ protected PgConnectOptions connectOptions() { @Override public void init() { codec = new PgCodec(useLayer7Proxy); + codec.setConnection(this); ChannelPipeline pipeline = socket.channelHandlerContext().pipeline(); pipeline.addBefore("handler", "codec", codec); super.init(); @@ -117,7 +118,9 @@ Future sendCancelRequestMessage(int processId, int secretKey) { @Override protected void handleMessage(Object msg) { super.handleMessage(msg); - if (msg instanceof Notification || msg instanceof TxFailedEvent || msg instanceof NoticeResponse) { + if (msg instanceof CopyOutEvent) { + ((CopyOutEvent) msg).dispatch(); + } else if (msg instanceof Notification || msg instanceof TxFailedEvent || msg instanceof NoticeResponse) { handleEvent(msg); } } diff --git a/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/CopyInHandler.java b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/CopyInHandler.java new file mode 100644 index 000000000..f798d0150 --- /dev/null +++ b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/CopyInHandler.java @@ -0,0 +1,16 @@ +/* + * Copyright (c) 2011-2026 Contributors to the Eclipse Foundation + * + * This program and the accompanying materials are made available under the + * terms of the Eclipse Public License 2.0 which is available at + * http://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0 + * which is available at https://www.apache.org/licenses/LICENSE-2.0. + * + * SPDX-License-Identifier: EPL-2.0 OR Apache-2.0 + */ + +package io.vertx.pgclient.impl.codec; + +interface CopyInHandler { + void handleCopyInResponse(int overall, short[] colFmts); +} diff --git a/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/CopyInStreamPgCommandMessage.java b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/CopyInStreamPgCommandMessage.java new file mode 100644 index 000000000..58a7fefe8 --- /dev/null +++ b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/CopyInStreamPgCommandMessage.java @@ -0,0 +1,241 @@ +/* + * Copyright (c) 2011-2026 Contributors to the Eclipse Foundation + * + * This program and the accompanying materials are made available under the + * terms of the Eclipse Public License 2.0 which is available at + * http://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0 + * which is available at https://www.apache.org/licenses/LICENSE-2.0. + * + * SPDX-License-Identifier: EPL-2.0 OR Apache-2.0 + */ + +package io.vertx.pgclient.impl.codec; + +import io.netty.buffer.ByteBuf; +import io.netty.channel.ChannelHandlerContext; +import io.netty.channel.ChannelInboundHandlerAdapter; +import io.netty.channel.ChannelPipeline; +import io.netty.channel.WriteBufferWaterMark; +import io.vertx.pgclient.impl.CopyInStreamCommand; +import io.vertx.pgclient.impl.CopyInStreamInternal; + +final class CopyInStreamPgCommandMessage + extends PgCommandMessage + implements CopyInHandler { + + private PgEncoder encoder; + + private final CopyInStreamInternal in; + + private int rowCount = -1; + + CopyInStreamPgCommandMessage(CopyInStreamCommand cmd) { + super(cmd); + this.in = cmd.in(); + } + + @Override + void encode(PgEncoder encoder) { + this.encoder = encoder; + + encoder.writeQuery(new QueryMessage(cmd.sql())); + + encoder.suspendCommandPipeline(); + } + + @Override + public void handleCopyInResponse(int overall, short[] colFmts) { + in.attachSink(new NettyCopyInSink(encoder)); + } + + @Override + public void handleCommandComplete(int updated) { + in.detachSinkIfAny(); + rowCount = updated; + result = null; + } + + @Override + public void handleErrorResponse(ErrorResponse err) { + in.detachSinkIfAny(); + // Recorded, not reported: the command completion fails the stream, after ReadyForQuery + failure = err.toException(); + } + + @Override + void handleReadyForQuery() { + encoder.resumeCommandPipeline(); + super.handleReadyForQuery(); + if (rowCount >= 0) { + // Held back until now so the connection is free when the application hears + in.completeFromServer(rowCount); + } + } + + private static final class NettyCopyInSink implements CopyInStreamInternal.Sink { + + private final PgEncoder encoder; + + private volatile ChannelHandlerContext ctx; + + private final String handlerName; + + private volatile Runnable writableCb; + + private volatile boolean detached; + + private boolean flushScheduled; + + private final WriteBufferWaterMark originalWatermark; + + private WriteBufferWaterMark copyWatermark; + + private final ChannelInboundHandlerAdapter writabilityHandler = new ChannelInboundHandlerAdapter() { + + @Override + public void handlerAdded(ChannelHandlerContext c) throws Exception { + ctx = c; + super.handlerAdded(c); + } + + @Override + public void channelWritabilityChanged(ChannelHandlerContext c) throws Exception { + if (!detached && c.channel().isWritable()) { + Runnable cb = writableCb; + if (cb != null) { + c.executor().execute(cb); + } + } + super.channelWritabilityChanged(c); + } + + @Override + public void channelReadComplete(ChannelHandlerContext c) throws Exception { + if (!detached && flushScheduled) { + doFlush(); + } + super.channelReadComplete(c); + } + }; + + NettyCopyInSink(PgEncoder encoder) { + this.encoder = encoder; + + ChannelHandlerContext encoderCtx = encoder.channelHandlerContext(); + if (encoderCtx == null) { + throw new IllegalStateException("PgEncoder.channelHandlerContext() returned null"); + } + + this.handlerName = "pg-copyin-writability@" + System.identityHashCode(this); + + ChannelPipeline p = encoderCtx.pipeline(); + if (p.get(handlerName) == null) { + p.addAfter("codec", handlerName, writabilityHandler); + } + + if (ctx == null) { + throw new IllegalStateException("COPY IN sink ctx was not initialized"); + } + + this.originalWatermark = ctx.channel().config().getWriteBufferWaterMark(); + } + + @Override + public boolean isWritable() { + return ctx.channel().isWritable(); + } + + @Override + public void onWritable(Runnable cb) { + this.writableCb = cb; + } + + @Override + public synchronized void setWatermarks(int maxBytes) { + if (detached) { + return; + } + try { + int hi = Math.max(64 * 1024, maxBytes); + int lo = Math.max(32 * 1024, hi / 2); + WriteBufferWaterMark watermark = new WriteBufferWaterMark(lo, hi); + ctx.channel().config().setWriteBufferWaterMark(watermark); + copyWatermark = watermark; + } catch (Throwable ignore) { + + } + } + + @Override + public void writeCopyData(ByteBuf buf) { + encoder.writeCopyData(buf); + flushAfterCopyDataWrite(); + } + + private void flushAfterCopyDataWrite() { + ChannelHandlerContext c = ctx; + if (!c.channel().isWritable() || c.channel().bytesBeforeUnwritable() == 0) { + doFlush(); + } else { + scheduleFlush(); + } + } + + @Override + public void writeCopyDone() { + encoder.writeCopyDone(); + doFlush(); + } + + @Override + public void writeCopyFail(String message) { + encoder.writeCopyFail(message); + doFlush(); + } + + @Override + public synchronized void detach() { + if (detached) { + return; + } + detached = true; + + try { + // Do not overwrite a watermark installed by another owner during COPY. + WriteBufferWaterMark currentWatermark = ctx.channel().config().getWriteBufferWaterMark(); + if (copyWatermark != null && currentWatermark == copyWatermark) { + ctx.channel().config().setWriteBufferWaterMark(originalWatermark); + } + } catch (Throwable ignore) { + } + + try { + ChannelPipeline p = ctx.pipeline(); + if (p.get(handlerName) != null) { + p.remove(handlerName); + } + } catch (Throwable ignore) { + } + } + + private void scheduleFlush() { + if (flushScheduled || detached) return; + flushScheduled = true; + + ChannelHandlerContext c = ctx; + c.executor().execute(() -> { + if (detached || !flushScheduled) { + return; + } + doFlush(); + }); + } + + private void doFlush() { + if (detached) return; + + flushScheduled = false; + ctx.flush(); + } + } +} diff --git a/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/CopyOutHandler.java b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/CopyOutHandler.java new file mode 100644 index 000000000..f2174d8c0 --- /dev/null +++ b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/CopyOutHandler.java @@ -0,0 +1,20 @@ +/* + * Copyright (c) 2011-2026 Contributors to the Eclipse Foundation + * + * This program and the accompanying materials are made available under the + * terms of the Eclipse Public License 2.0 which is available at + * http://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0 + * which is available at https://www.apache.org/licenses/LICENSE-2.0. + * + * SPDX-License-Identifier: EPL-2.0 OR Apache-2.0 + */ + +package io.vertx.pgclient.impl.codec; + +import io.netty.buffer.ByteBuf; + +interface CopyOutHandler { + void handleCopyOutResponse(int overall, short[] fmts); + void handleCopyData(ByteBuf data); + void handleCopyDone(); +} diff --git a/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/CopyOutStreamPgCommandMessage.java b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/CopyOutStreamPgCommandMessage.java new file mode 100644 index 000000000..7d3b64bb4 --- /dev/null +++ b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/CopyOutStreamPgCommandMessage.java @@ -0,0 +1,128 @@ +/* + * Copyright (c) 2011-2026 Contributors to the Eclipse Foundation + * + * This program and the accompanying materials are made available under the + * terms of the Eclipse Public License 2.0 which is available at + * http://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0 + * which is available at https://www.apache.org/licenses/LICENSE-2.0. + * + * SPDX-License-Identifier: EPL-2.0 OR Apache-2.0 + */ + +package io.vertx.pgclient.impl.codec; + +import io.netty.buffer.ByteBuf; +import io.vertx.core.buffer.Buffer; +import io.vertx.core.internal.buffer.BufferInternal; +import io.vertx.pgclient.impl.CopyOutEvent; +import io.vertx.pgclient.impl.CopyOutStreamCommand; +import io.vertx.pgclient.impl.CopyOutStreamImpl; + +final class CopyOutStreamPgCommandMessage + extends PgCommandMessage + implements CopyOutHandler { + + private final int aggregationThreshold; + private final CopyOutStreamImpl out; + private PgEncoder encoder; + private Buffer agg; + + CopyOutStreamPgCommandMessage(CopyOutStreamCommand cmd) { + super(cmd); + this.out = cmd.out(); + this.aggregationThreshold = cmd.options().getAggregationThreshold(); + this.agg = aggregationThreshold == 1 ? null : Buffer.buffer(aggregationThreshold); + } + + @Override + void encode(PgEncoder encoder) { + this.encoder = encoder; + encoder.writeQuery(new QueryMessage(cmd.sql())); + } + + @Override + public void handleCopyOutResponse(int overall, short[] colFmts) { + out.readyFromServer(); + } + + @Override + public void handleCopyData(ByteBuf data) { + if (failure != null || out.isDiscarding()) { + data.release(); + return; + } + + if (aggregationThreshold == 1) { + fireData(BufferInternal.safeBuffer(data)); + return; + } + + try { + int len = data.readableBytes(); + + if (agg.length() > 0 && agg.length() + len >= aggregationThreshold) { + flushAgg(); + } + + ((BufferInternal) agg).unwrap().writeBytes(data, len); + + if (agg.length() >= aggregationThreshold) { + flushAgg(); + } + } finally { + data.release(); + } + } + + @Override + public void handleCopyDone() { + // Nothing to do, CommandComplete carries the row count and ends the stream + } + + @Override + public void handleCommandComplete(int updated) { + if (out.isDiscarding()) { + result = null; + return; + } else if (failure != null) { + out.fail(failure); + } else { + flushAgg(); + fireEnd(updated); + } + + result = null; + } + + @Override + void handleReadyForQuery() { + super.handleReadyForQuery(); + // Same path as the data and end events, so it cannot overtake them + fireCompleted(); + } + + @Override + void handleErrorResponse(ErrorResponse errorResponse) { + // Recorded, not reported: the command completion fails the stream, after ReadyForQuery + failure = errorResponse.toException(); + } + + private void flushAgg() { + if (agg != null && agg.length() > 0) { + fireData(agg); + agg = Buffer.buffer(aggregationThreshold); + } + } + + private void fireData(Buffer data) { + encoder.channelHandlerContext().fireChannelRead(CopyOutEvent.data(out, data)); + } + + private void fireEnd(int rowCount) { + encoder.channelHandlerContext().fireChannelRead(CopyOutEvent.end(out, rowCount)); + } + + private void fireCompleted() { + encoder.channelHandlerContext().fireChannelRead(CopyOutEvent.completed(out)); + } +} diff --git a/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/PgCodec.java b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/PgCodec.java index 94fb47c6e..4b0eb996e 100644 --- a/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/PgCodec.java +++ b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/PgCodec.java @@ -16,12 +16,17 @@ */ package io.vertx.pgclient.impl.codec; +import io.netty.channel.ChannelHandlerContext; import io.netty.channel.CombinedChannelDuplexHandler; +import io.vertx.sqlclient.codec.SocketConnectionBase; import java.util.ArrayDeque; public class PgCodec extends CombinedChannelDuplexHandler { + private SocketConnectionBase connection; + private boolean commandPipelineSuspended; + private final ArrayDeque> inflight; private final PgDecoder decoder; private final PgEncoder encoder; @@ -33,6 +38,45 @@ public PgCodec(boolean useLayer7Proxy) { init(decoder, encoder); } + public void setConnection(SocketConnectionBase connection) { + this.connection = connection; + } + + @Override + public void channelInactive(ChannelHandlerContext ctx) throws Exception { + clearConnectionReference(); + super.channelInactive(ctx); + } + + @Override + public void handlerRemoved(ChannelHandlerContext ctx) throws Exception { + clearConnectionReference(); + super.handlerRemoved(ctx); + } + + private void clearConnectionReference() { + connection = null; + commandPipelineSuspended = false; + } + + void suspendCommandPipeline() { + if (!commandPipelineSuspended) { + commandPipelineSuspended = true; + if (connection != null) { + connection.suspendPipeline(); + } + } + } + + void resumeCommandPipeline() { + if (commandPipelineSuspended) { + commandPipelineSuspended = false; + if (connection != null) { + connection.resumePipeline(); + } + } + } + void add(PgCommandMessage codec) { codec.decoder = decoder; inflight.add(codec); diff --git a/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/PgCommandMessage.java b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/PgCommandMessage.java index 5ba457b4a..d936b629b 100644 --- a/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/PgCommandMessage.java +++ b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/PgCommandMessage.java @@ -20,6 +20,8 @@ import io.vertx.core.internal.logging.Logger; import io.vertx.core.internal.logging.LoggerFactory; import io.vertx.pgclient.PgException; +import io.vertx.pgclient.impl.CopyInStreamCommand; +import io.vertx.pgclient.impl.CopyOutStreamCommand; import io.vertx.sqlclient.codec.CommandMessage; import io.vertx.sqlclient.codec.CommandResponse; import io.vertx.sqlclient.spi.protocol.CloseConnectionCommand; @@ -57,6 +59,10 @@ public abstract class PgCommandMessage> extends Comm return new ClosePortalPgCommandMessage((CloseCursorCommand) cmd); } else if (cmd instanceof CloseStatementCommand) { return new CloseStatementPgCommandMessage((CloseStatementCommand) cmd); + } else if (cmd instanceof CopyOutStreamCommand) { + return new CopyOutStreamPgCommandMessage((CopyOutStreamCommand) cmd); + } else if (cmd instanceof CopyInStreamCommand) { + return new CopyInStreamPgCommandMessage((CopyInStreamCommand) cmd); } throw new AssertionError("Invalid command " + cmd); } diff --git a/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/PgDecoder.java b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/PgDecoder.java index da7dd7b2d..0d7c1a5ba 100644 --- a/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/PgDecoder.java +++ b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/PgDecoder.java @@ -117,6 +117,22 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) { decodeBindComplete(); break; } + case PgProtocolConstants.MESSAGE_TYPE_COPY_OUT_RESPONSE: { + decodeCopyOutResponse(in); + break; + } + case PgProtocolConstants.MESSAGE_TYPE_COPY_IN_RESPONSE: { + decodeCopyInResponse(in); + break; + } + case PgProtocolConstants.MESSAGE_TYPE_COPY_DATA: { + decodeCopyData(in); + break; + } + case PgProtocolConstants.MESSAGE_TYPE_COPY_DONE: { + decodeCopyDone(); + break; + } default: { decodeMessage(ctx, id, in); } @@ -195,6 +211,64 @@ private void decodePortalSuspended() { codec.peek().handlePortalSuspended(); } + private void decodeCopyOutResponse(ByteBuf in) { + PgCommandMessage msg = codec.peek(); + if (!(msg instanceof CopyOutHandler)) { + throw unexpectedCopyMessage("CopyOutResponse", msg); + } + + final int overall = in.readUnsignedByte(); + final int cols = in.readUnsignedShort(); + short[] fmts = new short[cols]; + for (int i = 0; i < cols; i++) { + fmts[i] = in.readShort(); + } + + ((CopyOutHandler) msg).handleCopyOutResponse(overall, fmts); + } + + private void decodeCopyData(ByteBuf in) { + PgCommandMessage msg = codec.peek(); + if (!(msg instanceof CopyOutHandler)) { + throw unexpectedCopyMessage("CopyData", msg); + } + + ByteBuf slice = in.readRetainedSlice(in.readableBytes()); + ((CopyOutHandler) msg).handleCopyData(slice); + } + + private void decodeCopyDone() { + PgCommandMessage msg = codec.peek(); + if (!(msg instanceof CopyOutHandler)) { + throw unexpectedCopyMessage("CopyDone", msg); + } + + ((CopyOutHandler) msg).handleCopyDone(); + } + + private void decodeCopyInResponse(ByteBuf in) { + PgCommandMessage msg = codec.peek(); + if (!(msg instanceof CopyInHandler)) { + throw unexpectedCopyMessage("CopyInResponse", msg); + } + + final int overall = in.readUnsignedByte(); + final int cols = in.readUnsignedShort(); + short[] fmts = new short[cols]; + for (int i = 0; i < cols; i++) { + fmts[i] = in.readShort(); + } + + ((CopyInHandler) msg).handleCopyInResponse(overall, fmts); + } + + private RuntimeException unexpectedCopyMessage(String message, PgCommandMessage msg) { + return new IllegalStateException( + "Unexpected PostgreSQL " + message + " message for " + + (msg != null ? msg.getClass().getName() : "no current command") + ); + } + private void decodeCommandComplete(ByteBuf in) { int updated = processor.parse(in); codec.peek().handleCommandComplete(updated); diff --git a/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/PgEncoder.java b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/PgEncoder.java index c47b1723a..4cbae6373 100644 --- a/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/PgEncoder.java +++ b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/PgEncoder.java @@ -29,6 +29,7 @@ import io.vertx.sqlclient.internal.RowDescriptorBase; import io.vertx.sqlclient.spi.protocol.CloseConnectionCommand; +import java.nio.charset.StandardCharsets; import java.util.*; import static io.vertx.pgclient.impl.util.Util.writeCString; @@ -51,6 +52,9 @@ final class PgEncoder extends ChannelOutboundHandlerAdapter { private static final byte EXECUTE = 'E'; private static final byte CLOSE = 'C'; private static final byte SYNC = 'S'; + private static final byte COPY_DATA = 'd'; + private static final byte COPY_DONE = 'c'; + private static final byte COPY_FAIL = 'f'; private final PgCodec codec; final boolean useLayer7Proxy; @@ -66,6 +70,14 @@ final class PgEncoder extends ChannelOutboundHandlerAdapter { this.codec = codec; } + void suspendCommandPipeline() { + codec.suspendCommandPipeline(); + } + + void resumeCommandPipeline() { + codec.resumeCommandPipeline(); + } + private void enqueueMessage(Object msg, int estimate) { pendingMessages.add(msg); capacityEstimate += estimate; @@ -643,6 +655,30 @@ void writeBind(BindMessage bind, String portal, Tuple paramValues) { enqueueMessage(bind, portal, paramValues, estimateBind(bind, portal, paramValues)); } + void writeCopyData(ByteBuf payload) { + ByteBuf header = ctx.alloc().buffer(5, 5); + header.writeByte(COPY_DATA); + header.writeInt(payload.readableBytes() + 4); + ctx.write(Unpooled.wrappedBuffer(header, payload), ctx.voidPromise()); + } + + void writeCopyDone() { + ByteBuf msg = ctx.alloc().buffer(5, 5); + msg.writeByte(COPY_DONE); + msg.writeInt(4); + ctx.write(msg, ctx.voidPromise()); + } + + void writeCopyFail(String message) { + byte[] msgBytes = message != null ? message.getBytes(StandardCharsets.UTF_8) : new byte[0]; + ByteBuf msg = ctx.alloc().buffer(1 + 4 + msgBytes.length + 1); + msg.writeByte(COPY_FAIL); + msg.writeInt(4 + msgBytes.length + 1); + msg.writeBytes(msgBytes); + msg.writeByte(0); + ctx.write(msg, ctx.voidPromise()); + } + byte[] nextStatementName() { return psSeq.next(); } diff --git a/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/PgProtocolConstants.java b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/PgProtocolConstants.java index bc8fa9a0d..60cfa0a13 100644 --- a/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/PgProtocolConstants.java +++ b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/PgProtocolConstants.java @@ -71,4 +71,9 @@ public class PgProtocolConstants { public static final byte MESSAGE_TYPE_FUNCTION_RESULT = 'V'; public static final byte MESSAGE_TYPE_SSL_YES = 'S'; public static final byte MESSAGE_TYPE_SSL_NO = 'N'; + + public static final byte MESSAGE_TYPE_COPY_DATA = 'd'; + public static final byte MESSAGE_TYPE_COPY_DONE = 'c'; + public static final byte MESSAGE_TYPE_COPY_IN_RESPONSE = 'G'; + public static final byte MESSAGE_TYPE_COPY_OUT_RESPONSE = 'H'; } diff --git a/vertx-pg-client/src/test/java/io/vertx/tests/pgclient/CopyStreamTest.java b/vertx-pg-client/src/test/java/io/vertx/tests/pgclient/CopyStreamTest.java new file mode 100644 index 000000000..ed31b73d3 --- /dev/null +++ b/vertx-pg-client/src/test/java/io/vertx/tests/pgclient/CopyStreamTest.java @@ -0,0 +1,1075 @@ +/* + * Copyright (c) 2011-2026 Contributors to the Eclipse Foundation + * + * This program and the accompanying materials are made available under the + * terms of the Eclipse Public License 2.0 which is available at + * http://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0 + * which is available at https://www.apache.org/licenses/LICENSE-2.0. + * + * SPDX-License-Identifier: EPL-2.0 OR Apache-2.0 + */ + +package io.vertx.tests.pgclient; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.CompositeByteBuf; +import io.vertx.core.Future; +import io.vertx.core.Handler; +import io.vertx.core.Vertx; +import io.vertx.core.buffer.Buffer; +import io.vertx.core.internal.ContextInternal; +import io.vertx.core.internal.buffer.BufferInternal; +import io.vertx.core.json.JsonObject; +import io.vertx.core.streams.ReadStream; +import io.vertx.ext.unit.Async; +import io.vertx.ext.unit.TestContext; +import io.vertx.ext.unit.junit.VertxUnitRunner; +import io.vertx.pgclient.PgCopyInOptions; +import io.vertx.pgclient.PgCopyOutOptions; +import io.vertx.pgclient.impl.CopyInStreamImpl; +import io.vertx.pgclient.impl.CopyInStreamInternal; +import io.vertx.pgclient.impl.CopyOutStreamImpl; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotSame; +import static org.junit.Assert.fail; + +/** + * Unit tests for the COPY streams and their options, with no database involved. + */ +@RunWith(VertxUnitRunner.class) +public class CopyStreamTest { + + + private Vertx vertx; + + @Before + public void setUp() { + vertx = Vertx.vertx(); + } + + @After + public void tearDown(TestContext ctx) { + vertx.close().onComplete(ctx.asyncAssertSuccess()); + } + + @Test + public void testSequentialComposedWritesProgress(TestContext ctx) { + assertSequentialComposedWrites(ctx); + } + + @Test + public void testSameTurnWritesCoalesceAtChunkBoundaryWithoutLosingOversizedWrite(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(ignored -> { + CopyInStreamImpl stream = coalescing(4); + TestSink sink = new TestSink(); + completeFromDone(stream, sink, 3); + + attach(stream, sink).onComplete(ctx.asyncAssertSuccess(v -> { + Future w1 = stream.write(buffer("ab")); + Future w2 = stream.write(buffer("cd")); + Future w3 = stream.write(buffer("efghij")); + Future end = stream.end(); + + Future.all(w1, w2, w3).compose(x -> end).onComplete(ctx.asyncAssertSuccess(x -> { + ctx.assertEquals(Arrays.asList("data:abcd", "data:efghij", "done"), sink.events); + async.complete(); + })); + })); + }); + } + + @Test + public void testEmptyCoalescingWriteIsNoOpAndDoesNotRetainBuffer(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(ignored -> { + CopyInStreamImpl stream = coalescing(1024); + TestSink sink = new TestSink(); + completeFromDone(stream, sink, 0); + + attach(stream, sink).onComplete(ctx.asyncAssertSuccess(v -> { + Buffer empty = Buffer.buffer(); + ByteBuf byteBuf = ((BufferInternal) empty).getByteBuf(); + int refCnt = byteBuf.refCnt(); + Future write = stream.write(empty); + Future end = stream.end(); + + write.compose(x -> end).onComplete(ctx.asyncAssertSuccess(x -> { + ctx.assertEquals(refCnt, byteBuf.refCnt()); + ctx.assertEquals(Arrays.asList("done"), sink.events); + async.complete(); + })); + })); + }); + } + + @Test + public void testUnwritableSinkKeepsDataBeforeCopyDoneAndFiresSingleDrain(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(ignored -> { + CopyInStreamImpl stream = coalescing(64); + TestSink sink = new TestSink(); + sink.writable = false; + AtomicInteger drains = new AtomicInteger(); + completeFromDone(stream, sink, 1); + + attach(stream, sink).onComplete(ctx.asyncAssertSuccess(v -> { + stream.setWriteQueueMaxSize(1); + stream.drainHandler(x -> drains.incrementAndGet()); + Future write = stream.write(buffer("x")); + Future end = stream.end(); + + vertx.runOnContext(next -> { + ctx.assertTrue(stream.writeQueueFull()); + ctx.assertFalse(write.isComplete()); + ctx.assertFalse(end.isComplete()); + ctx.assertTrue(sink.events.isEmpty()); + + sink.makeWritable(); + end.onComplete(ctx.asyncAssertSuccess(x -> { + ctx.assertEquals(Arrays.asList("data:x", "done"), sink.events); + ctx.assertEquals(1, drains.get()); + ctx.assertFalse(stream.writeQueueFull()); + ctx.assertTrue(write.succeeded()); + async.complete(); + })); + }); + })); + }); + } + + @Test + public void testAbortReleasesBufferedDataAndRejectsFurtherOperations(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(ignored -> { + CopyInStreamImpl stream = coalescing(1024); + TestSink sink = new TestSink(); + sink.writable = false; + stream.exceptionHandler(t -> { }); + + attach(stream, sink).onComplete(ctx.asyncAssertSuccess(v -> { + Buffer data = buffer("buffered"); + ByteBuf byteBuf = ((BufferInternal) data).getByteBuf(); + int refCnt = byteBuf.refCnt(); + Future write = stream.write(data); + + // PostgreSQL answers CopyFail with an error, and that error is the outcome of the COPY + IllegalStateException serverError = new IllegalStateException("COPY from stdin failed: rollback"); + + stream.abort("rollback").compose(x -> { + // aborting stops the stream, but the outcome is still the server's to report + ctx.assertFalse(stream.completion().isComplete()); + stream.failFromServer(serverError); + return Future.all( + expectFailure(write), + expectFailure(stream.completion()), + expectFailure(stream.write(buffer("late"))), + expectFailure(stream.abort("again"))); + }).onComplete(ctx.asyncAssertSuccess(x -> { + ctx.assertEquals(serverError, stream.completion().cause()); + ctx.assertEquals(refCnt, byteBuf.refCnt()); + ctx.assertEquals(Arrays.asList("fail:rollback"), sink.events); + ctx.assertEquals(1, sink.copyFailCalls); + async.complete(); + })); + })); + }); + } + + @Test + public void testServerFailureReleasesUnwrittenDataWithoutSendingCopyFail(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(ignored -> { + CopyInStreamImpl stream = coalescing(1024); + TestSink sink = new TestSink(); + sink.writable = false; + stream.exceptionHandler(t -> { }); + + attach(stream, sink).onComplete(ctx.asyncAssertSuccess(v -> { + Buffer data = buffer("buffered"); + ByteBuf byteBuf = ((BufferInternal) data).getByteBuf(); + int refCnt = byteBuf.refCnt(); + Future write = stream.write(data); + stream.failFromServer(new IllegalStateException("server rejected COPY")); + + Future.all(expectFailure(write), expectFailure(stream.completion())).onComplete(ctx.asyncAssertSuccess(x -> { + ctx.assertEquals(refCnt, byteBuf.refCnt()); + ctx.assertTrue(sink.events.isEmpty()); + ctx.assertEquals(0, sink.copyFailCalls); + async.complete(); + })); + })); + }); + } + + @Test + public void testEndIsIdempotentAndCopyDoneIsLast(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(ignored -> { + CopyInStreamImpl stream = coalescing(1024); + TestSink sink = new TestSink(); + completeFromDone(stream, sink, 0); + + attach(stream, sink).onComplete(ctx.asyncAssertSuccess(v -> { + Future end1 = stream.end(); + Future end2 = stream.end(); + + Future.all(end1, end2).compose(x -> Future.all( + expectFailure(stream.write(buffer("late"))), + expectFailure(stream.abort("late"))) + ).onComplete(ctx.asyncAssertSuccess(x -> { + ctx.assertEquals(Arrays.asList("done"), sink.events); + ctx.assertEquals(1, sink.copyDoneCalls); + ctx.assertEquals(0, sink.copyFailCalls); + async.complete(); + })); + })); + }); + } + + @Test + public void testSynchronousSinkFailureReleasesFrame(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(ignored -> { + CopyInStreamImpl stream = coalescing(1024); + TestSink sink = new TestSink(); + sink.throwOnWrite = new IllegalStateException("write threw"); + stream.exceptionHandler(t -> { }); + + attach(stream, sink).onComplete(ctx.asyncAssertSuccess(v -> { + Buffer data = buffer("payload"); + ByteBuf byteBuf = ((BufferInternal) data).getByteBuf(); + int refCnt = byteBuf.refCnt(); + + Future.all(expectFailure(stream.write(data)), expectFailure(stream.completion())) + .onComplete(ctx.asyncAssertSuccess(x -> { + ctx.assertEquals(refCnt, byteBuf.refCnt()); + ctx.assertEquals(Arrays.asList("fail:write threw"), sink.events); + ctx.assertEquals(1, sink.copyFailCalls); + async.complete(); + })); + })); + }); + } + + private void assertSequentialComposedWrites(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(ignored -> { + CopyInStreamImpl stream = coalescing(1024); + TestSink sink = new TestSink(); + completeFromDone(stream, sink, 3); + + attach(stream, sink).onComplete(ctx.asyncAssertSuccess(v -> stream.write(buffer("a")) + .compose(x -> stream.write(buffer("b"))) + .compose(x -> stream.write(buffer("c"))) + .compose(x -> stream.end()) + .onComplete(ctx.asyncAssertSuccess(x -> { + ctx.assertEquals(Arrays.asList("data:a", "data:b", "data:c", "done"), sink.events); + async.complete(); + })))); + }); + } + + private CopyInStreamImpl coalescing(int chunkSize) { + return new CopyInStreamImpl(currentContext(), new PgCopyInOptions() + .setChunkSize(chunkSize)); + } + + private static ContextInternal currentContext() { + return (ContextInternal) Vertx.currentContext(); + } + + private static Future attach(CopyInStreamInternal stream, TestSink sink) { + stream.attachSink(sink); + return stream.readyFuture(); + } + + private static void completeFromDone(CopyInStreamInternal stream, TestSink sink, int rows) { + sink.doneHandler = () -> stream.completeFromServer(rows); + } + + private static Future expectFailure(Future future) { + return future.compose(v -> Future.failedFuture("expected failure"), err -> Future.succeededFuture()); + } + + private static Buffer buffer(String value) { + return Buffer.buffer(value, StandardCharsets.UTF_8.name()); + } + + private static final class TestSink implements CopyInStreamInternal.Sink { + + final List events = new ArrayList<>(); + + boolean writable = true; + RuntimeException throwOnWrite; + Runnable writableHandler; + Runnable doneHandler; + int copyDoneCalls; + int copyFailCalls; + int detachCalls; + int watermark; + Integer lastFrameComponents; + String lastCopyFailMessage; + + @Override + public boolean isWritable() { + return writable; + } + + @Override + public void onWritable(Runnable cb) { + writableHandler = cb; + } + + @Override + public void setWatermarks(int maxBytes) { + watermark = maxBytes; + } + + @Override + public void writeCopyData(ByteBuf buf) { + if (throwOnWrite != null) { + throw throwOnWrite; + } + lastFrameComponents = buf instanceof CompositeByteBuf ? ((CompositeByteBuf) buf).numComponents() : 1; + record(buf); + } + + @Override + public void writeCopyDone() { + copyDoneCalls++; + events.add("done"); + if (doneHandler != null) { + doneHandler.run(); + } + } + + @Override + public void writeCopyFail(String message) { + copyFailCalls++; + lastCopyFailMessage = message; + events.add("fail:" + message); + } + + @Override + public void detach() { + detachCalls++; + } + + void makeWritable() { + writable = true; + if (writableHandler != null) { + writableHandler.run(); + } + } + + private void record(ByteBuf buf) { + byte[] bytes = new byte[buf.readableBytes()]; + buf.getBytes(buf.readerIndex(), bytes); + events.add("data:" + new String(bytes, StandardCharsets.UTF_8)); + buf.release(); + } + } + + // ---------------------------------------------------------------- COPY OUT + + @Test + public void testReentrantFetchDrainsIteratively(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(v -> { + TestReadStream upstream = new TestReadStream(); + CopyOutStreamImpl out = new CopyOutStreamImpl((ContextInternal) Vertx.currentContext(), upstream); + + out.pause(); + ctx.assertEquals(1, upstream.pauseCalls); + for (int i = 0; i < 16; i++) { + out.emit(Buffer.buffer(Integer.toString(i))); + } + + ctx.assertEquals(1, upstream.pauseCalls); + ctx.assertEquals(0, upstream.resumeCalls); + + AtomicInteger delivered = new AtomicInteger(); + AtomicInteger ended = new AtomicInteger(); + AtomicInteger handlerDepth = new AtomicInteger(); + AtomicInteger maxHandlerDepth = new AtomicInteger(); + out.handler(buffer -> { + int depth = handlerDepth.incrementAndGet(); + maxHandlerDepth.accumulateAndGet(depth, Math::max); + if (delivered.incrementAndGet() < 16) { + out.fetch(1); + } + handlerDepth.decrementAndGet(); + }); + out.endHandler(ignored -> ended.incrementAndGet()); + out.end(16); + // CommandComplete ends the stream, ReadyForQuery resolves completion() + ctx.assertFalse(out.completion().succeeded()); + out.commandCompleted(); + ctx.assertEquals(0, ended.get()); + out.fetch(1); + + vertx.runOnContext(ignored -> { + ctx.assertTrue(out.completion().succeeded()); + ctx.assertEquals(16, delivered.get()); + ctx.assertEquals(1, ended.get()); + ctx.assertEquals(1, maxHandlerDepth.get()); + ctx.assertEquals(1, upstream.pauseCalls); + ctx.assertEquals(1, upstream.resumeCalls); + async.complete(); + }); + }); + } + + @Test + public void testDiscardReleasesPausedUpstream(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(v -> { + TestReadStream upstream = new TestReadStream(); + CopyOutStreamImpl out = new CopyOutStreamImpl((ContextInternal) Vertx.currentContext(), upstream); + + out.emit(Buffer.buffer("data")); + ctx.assertEquals(1, upstream.pauseCalls); + + out.discard(new IllegalStateException("discarded")); + + ctx.assertEquals(1, upstream.resumeCalls); + ctx.assertTrue(out.completion().failed()); + async.complete(); + }); + } + + @Test + public void testInitiallyFlowingHandlerDoesNotTouchUpstream(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(v -> { + TestReadStream upstream = new TestReadStream(); + CopyOutStreamImpl out = new CopyOutStreamImpl((ContextInternal) Vertx.currentContext(), upstream); + AtomicInteger delivered = new AtomicInteger(); + + out.handler(buffer -> delivered.incrementAndGet()); + out.emit(Buffer.buffer("a")); + out.emit(Buffer.buffer("b")); + + ctx.assertEquals(2, delivered.get()); + ctx.assertEquals(0, upstream.pauseCalls); + ctx.assertEquals(0, upstream.resumeCalls); + async.complete(); + }); + } + + @Test + public void testQueuedDataPausesUpstreamImmediately(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(v -> { + TestReadStream upstream = new TestReadStream(); + CopyOutStreamImpl out = new CopyOutStreamImpl((ContextInternal) Vertx.currentContext(), upstream); + AtomicInteger delivered = new AtomicInteger(); + + out.emit(Buffer.buffer("a")); + out.emit(Buffer.buffer("b")); // Simulates an event already decoded when pause took effect. + + ctx.assertEquals(1, upstream.pauseCalls); + out.handler(buffer -> delivered.incrementAndGet()); + ctx.assertEquals(2, delivered.get()); + ctx.assertEquals(1, upstream.resumeCalls); + async.complete(); + }); + } + + @Test + public void testExhaustedDemandPausesUpstreamImmediately(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(v -> { + TestReadStream upstream = new TestReadStream(); + CopyOutStreamImpl out = new CopyOutStreamImpl((ContextInternal) Vertx.currentContext(), upstream); + AtomicInteger delivered = new AtomicInteger(); + + out.handler(buffer -> delivered.incrementAndGet()); + out.pause(); + out.fetch(1); + + ctx.assertEquals(1, upstream.pauseCalls); + ctx.assertEquals(1, upstream.resumeCalls); + out.emit(Buffer.buffer("a")); + ctx.assertEquals(1, delivered.get()); + ctx.assertEquals(2, upstream.pauseCalls); + async.complete(); + }); + } + + @Test + public void testEndWaitsForQueuedDataAndIsEmittedOnce(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(v -> { + TestReadStream upstream = new TestReadStream(); + CopyOutStreamImpl out = new CopyOutStreamImpl((ContextInternal) Vertx.currentContext(), upstream); + AtomicInteger delivered = new AtomicInteger(); + AtomicInteger ended = new AtomicInteger(); + + out.pause(); + out.handler(buffer -> delivered.incrementAndGet()); + out.endHandler(ignored -> ended.incrementAndGet()); + out.emit(Buffer.buffer("a")); + out.emit(Buffer.buffer("b")); + out.end(2); + out.commandCompleted(); + + ctx.assertEquals(0, ended.get()); + out.fetch(1); + ctx.assertEquals(1, delivered.get()); + ctx.assertEquals(0, ended.get()); + out.fetch(1); + out.resume(); + out.fetch(1); + out.endHandler(ignored -> ended.incrementAndGet()); + + vertx.runOnContext(ignored -> { + ctx.assertTrue(out.completion().succeeded()); + ctx.assertEquals(2, delivered.get()); + ctx.assertEquals(1, ended.get()); + async.complete(); + }); + }); + } + + @Test + public void testHandlerCanClearItselfDuringDrain(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(v -> { + TestReadStream upstream = new TestReadStream(); + CopyOutStreamImpl out = new CopyOutStreamImpl((ContextInternal) Vertx.currentContext(), upstream); + AtomicInteger delivered = new AtomicInteger(); + + out.emit(Buffer.buffer("a")); + out.emit(Buffer.buffer("b")); + out.emit(Buffer.buffer("c")); + out.handler(buffer -> { + delivered.incrementAndGet(); + out.handler(null); + }); + + ctx.assertEquals(1, delivered.get()); + out.handler(buffer -> delivered.incrementAndGet()); + ctx.assertEquals(3, delivered.get()); + async.complete(); + }); + } + + // ---------------------------------------------------------------- edge cases + + @Test + public void testExceptionHandlerRegisteredAfterFailureIsStillNotified(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(v -> { + CopyOutStreamImpl out = copyOut(); + RuntimeException boom = new RuntimeException("boom"); + out.fail(boom); + + // handler installed after the failure already happened + out.exceptionHandler(err -> { + ctx.assertEquals(boom, err); + async.complete(); + }); + }); + } + + @Test + public void testFetchRejectsNegativeAmount(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(v -> { + CopyOutStreamImpl out = copyOut(); + try { + out.fetch(-1); + ctx.fail("expected an IllegalArgumentException"); + } catch (IllegalArgumentException expected) { + } + // zero is a legal no-op + out.fetch(0); + async.complete(); + }); + } + + @Test + public void testEmitAfterEndIsIgnored(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(v -> { + CopyOutStreamImpl out = copyOut(); + AtomicInteger delivered = new AtomicInteger(); + out.handler(b -> delivered.incrementAndGet()); + out.end(0); + out.emit(Buffer.buffer("late")); + + vertx.runOnContext(x -> { + ctx.assertEquals(0, delivered.get()); + async.complete(); + }); + }); + } + + @Test + public void testEmitOfNullIsIgnored(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(v -> { + CopyOutStreamImpl out = copyOut(); + AtomicInteger delivered = new AtomicInteger(); + out.handler(b -> delivered.incrementAndGet()); + out.emit(null); + + vertx.runOnContext(x -> { + ctx.assertEquals(0, delivered.get()); + async.complete(); + }); + }); + } + + @Test + public void testFirstFailureWins(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(v -> { + CopyOutStreamImpl out = copyOut(); + RuntimeException first = new RuntimeException("first"); + out.fail(first); + out.fail(new RuntimeException("second")); + + out.completion().onComplete(ar -> { + ctx.assertTrue(ar.failed()); + ctx.assertEquals(first, ar.cause()); + async.complete(); + }); + }); + } + + @Test + public void testFailureAfterTheDataDrainedStillSettlesCompletion(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(v -> { + CopyOutStreamImpl out = copyOut(); + out.handler(b -> { + }); + // CommandComplete ends the stream, the connection can still die before ReadyForQuery + out.end(3); + ctx.assertFalse(out.completion().isComplete()); + + RuntimeException boom = new RuntimeException("connection closed"); + out.fail(boom); + + out.completion().onComplete(ar -> { + ctx.assertTrue(ar.failed()); + ctx.assertEquals(boom, ar.cause()); + async.complete(); + }); + }); + } + + @Test + public void testDiscardIsIdempotentAndFailsCompletion(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(v -> { + TestReadStream upstream = new TestReadStream(); + CopyOutStreamImpl out = new CopyOutStreamImpl(currentContext(), upstream); + out.emit(Buffer.buffer("queued")); + + out.discard(new IllegalStateException("first discard")); + out.discard(new IllegalStateException("second discard")); + + ctx.assertTrue(out.isDiscarding()); + ctx.assertTrue(out.completion().failed()); + ctx.assertEquals("first discard", out.completion().cause().getMessage()); + // the upstream is released exactly once + ctx.assertEquals(1, upstream.resumeCalls); + async.complete(); + }); + } + + @Test + public void testReadyFutureCompletesWhenServerEntersCopyMode(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(v -> { + CopyOutStreamImpl out = copyOut(); + ctx.assertFalse(out.readyFuture().isComplete()); + out.readyFromServer(); + out.readyFuture().onComplete(ctx.asyncAssertSuccess(stream -> { + ctx.assertEquals(out, stream); + async.complete(); + })); + }); + } + + @Test + public void testCompletionCarriesTheRowCountFromEnd(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(v -> { + CopyOutStreamImpl out = copyOut(); + out.handler(b -> { + }); + out.end(7); + ctx.assertFalse(out.completion().isComplete()); + out.commandCompleted(); + ctx.assertTrue(out.completion().succeeded()); + ctx.assertEquals(7, out.completion().result()); + async.complete(); + }); + } + + @Test + public void testEndedStreamNeverPausesUpstreamAgain(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(v -> { + TestReadStream upstream = new TestReadStream(); + CopyOutStreamImpl out = new CopyOutStreamImpl(currentContext(), upstream); + out.handler(b -> { + }); + out.end(0); + + int pausesAfterEnd = upstream.pauseCalls; + out.pause(); + out.fetch(0); + ctx.assertEquals(pausesAfterEnd, upstream.pauseCalls); + async.complete(); + }); + } + + @Test + public void testAbortBeforeTheSinkIsAttachedSendsCopyFailOnAttach(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(v -> { + CopyInStreamImpl in = copyIn(); + TestSink sink = new TestSink(); + + in.abort("early abort").onComplete(ctx.asyncAssertSuccess(x -> { + in.attachSink(sink); + vertx.runOnContext(y -> { + ctx.assertEquals(1, sink.copyFailCalls); + ctx.assertEquals("early abort", sink.lastCopyFailMessage); + // the sink is released again straight away, the copy never starts + ctx.assertEquals(1, sink.detachCalls); + ctx.assertTrue(in.readyFuture().failed()); + async.complete(); + }); + })); + }); + } + + @Test + public void testAbortAfterEndIsRejected(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(v -> { + CopyInStreamImpl in = copyIn(); + TestSink sink = new TestSink(); + in.attachSink(sink); + + in.readyFuture().onComplete(ctx.asyncAssertSuccess(x -> { + in.end(); + in.abort("too late").onComplete(ar -> { + ctx.assertTrue(ar.failed()); + ctx.assertTrue(ar.cause() instanceof IllegalStateException, "was " + ar.cause()); + // no CopyFail was sent, the copy is already finishing + ctx.assertEquals(0, sink.copyFailCalls); + async.complete(); + }); + })); + }); + } + + @Test + public void testWriteQueueMaxSizeIsClampedAndReachesTheSink(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(v -> { + CopyInStreamImpl in = copyIn(); + TestSink sink = new TestSink(); + in.attachSink(sink); + + in.readyFuture().onComplete(ctx.asyncAssertSuccess(x -> { + in.setWriteQueueMaxSize(4096); + ctx.assertEquals(4096, sink.watermark); + + // anything below one byte is clamped rather than rejected + in.setWriteQueueMaxSize(0); + ctx.assertEquals(1, sink.watermark); + in.setWriteQueueMaxSize(-99); + ctx.assertEquals(1, sink.watermark); + async.complete(); + })); + }); + } + + @Test + public void testDetachSinkWithoutASinkIsANoOp(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(v -> { + CopyInStreamImpl in = copyIn(); + in.detachSinkIfAny(); + vertx.runOnContext(x -> { + ctx.assertFalse(in.completion().isComplete()); + async.complete(); + }); + }); + } + + @Test + public void testCompleteFromServerResolvesCompletion(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(v -> { + CopyInStreamImpl in = copyIn(); + in.completeFromServer(12); + in.completion().onComplete(ctx.asyncAssertSuccess(rows -> { + ctx.assertEquals(12, rows); + async.complete(); + })); + }); + } + + @Test + public void testWriteOfNullIsRejected(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(v -> { + CopyInStreamImpl in = copyIn(); + TestSink sink = new TestSink(); + in.attachSink(sink); + + in.readyFuture().onComplete(ctx.asyncAssertSuccess(x -> in.write(null).onComplete(ar -> { + ctx.assertTrue(ar.failed()); + ctx.assertTrue(ar.cause() instanceof NullPointerException, "was " + ar.cause()); + async.complete(); + }))); + }); + } + + @Test + public void testServerFailureAfterAbortReplacesTheCompletionCause(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(v -> { + CopyInStreamImpl in = copyIn(); + TestSink sink = new TestSink(); + in.attachSink(sink); + + in.readyFuture().onComplete(ctx.asyncAssertSuccess(x -> in.abort("client abort").onComplete(ar -> { + // abort alone does not decide the outcome + ctx.assertFalse(in.completion().isComplete()); + + RuntimeException serverError = new RuntimeException("server said no"); + in.failFromServer(serverError); + + in.completion().onComplete(done -> { + ctx.assertTrue(done.failed()); + ctx.assertEquals(serverError, done.cause()); + async.complete(); + }); + }))); + }); + } + + @Test + public void testCoalescingKeepsSmallWritesUncopied(TestContext ctx) { + Async async = ctx.async(); + vertx.runOnContext(v -> { + CopyInStreamImpl in = new CopyInStreamImpl(currentContext(), new PgCopyInOptions() + .setChunkSize(1024 * 1024)); + TestSink sink = new TestSink(); + in.attachSink(sink); + + in.readyFuture().onComplete(ctx.asyncAssertSuccess(x -> { + int writes = 200; + for (int i = 0; i < writes; i++) { + in.write(Buffer.buffer("row-" + i + "\n")); + } + + vertx.runOnContext(y -> { + // Netty consolidates a CompositeByteBuf once it grows past its component limit, and + // consolidating copies every byte accumulated so far. The default limit is 16, which + // would turn coalescing many small writes into repeated whole chunk copies. Seeing one + // component per write proves the payload was assembled without copying. + ctx.assertNotNull(sink.lastFrameComponents, "no CopyData frame reached the sink"); + ctx.assertEquals(writes, sink.lastFrameComponents, + "the chunk was consolidated, so small writes are being copied repeatedly"); + in.abort("done").onComplete(z -> async.complete()); + }); + })); + }); + } + + // ---------------------------------------------------------------- options + + @Test + public void testCopyInDefaults() { + PgCopyInOptions options = new PgCopyInOptions(); + assertEquals(PgCopyInOptions.DEFAULT_CHUNK_SIZE, options.getChunkSize()); + assertEquals(256 * 1024, options.getChunkSize()); + } + + @Test + public void testCopyInSettersAreFluentAndStick() { + PgCopyInOptions options = new PgCopyInOptions(); + assertEquals(options, options.setChunkSize(4096)); + + assertEquals(4096, options.getChunkSize()); + } + + @Test + public void testCopyInCopyConstructor() { + PgCopyInOptions original = new PgCopyInOptions() + .setChunkSize(1234); + + PgCopyInOptions copy = new PgCopyInOptions(original); + assertNotSame(original, copy); + assertEquals(original.getChunkSize(), copy.getChunkSize()); + + // the copy is independent + copy.setChunkSize(9999); + assertEquals(1234, original.getChunkSize()); + } + + @Test + public void testCopyInJsonRoundTrip() { + PgCopyInOptions original = new PgCopyInOptions() + .setChunkSize(8192); + + JsonObject json = original.toJson(); + assertEquals(8192, (int) json.getInteger("chunkSize")); + + PgCopyInOptions restored = new PgCopyInOptions(json); + assertEquals(original.getChunkSize(), restored.getChunkSize()); + } + + @Test + public void testCopyInFromEmptyJsonKeepsDefaults() { + PgCopyInOptions options = new PgCopyInOptions(new JsonObject()); + assertEquals(PgCopyInOptions.DEFAULT_CHUNK_SIZE, options.getChunkSize()); + } + + @Test + public void testCopyInFromJsonWithOnlyChunkSize() { + PgCopyInOptions options = new PgCopyInOptions(new JsonObject().put("chunkSize", 4096)); + assertEquals(4096, options.getChunkSize()); + } + + @Test + public void testCopyInRejectsInvalidChunkSize() { + expectIllegalArgument(() -> new PgCopyInOptions().setChunkSize(0)); + expectIllegalArgument(() -> new PgCopyInOptions().setChunkSize(-1)); + expectIllegalArgument(() -> new PgCopyInOptions().setChunkSize(Integer.MIN_VALUE)); + // one byte is degenerate but legal + assertEquals(1, new PgCopyInOptions().setChunkSize(1).getChunkSize()); + } + + @Test + public void testCopyOutDefaults() { + PgCopyOutOptions options = new PgCopyOutOptions(); + assertEquals(PgCopyOutOptions.DEFAULT_AGGREGATION_THRESHOLD, options.getAggregationThreshold()); + assertEquals(1, options.getAggregationThreshold()); + } + + @Test + public void testCopyOutSetterIsFluentAndSticks() { + PgCopyOutOptions options = new PgCopyOutOptions(); + assertEquals(options, options.setAggregationThreshold(64 * 1024)); + assertEquals(64 * 1024, options.getAggregationThreshold()); + } + + @Test + public void testCopyOutCopyConstructor() { + PgCopyOutOptions original = new PgCopyOutOptions().setAggregationThreshold(4096); + PgCopyOutOptions copy = new PgCopyOutOptions(original); + assertNotSame(original, copy); + assertEquals(4096, copy.getAggregationThreshold()); + + copy.setAggregationThreshold(1); + assertEquals(4096, original.getAggregationThreshold()); + } + + @Test + public void testCopyOutJsonRoundTrip() { + PgCopyOutOptions original = new PgCopyOutOptions().setAggregationThreshold(2048); + JsonObject json = original.toJson(); + assertEquals(2048, (int) json.getInteger("aggregationThreshold")); + + PgCopyOutOptions restored = new PgCopyOutOptions(json); + assertEquals(original.getAggregationThreshold(), restored.getAggregationThreshold()); + } + + @Test + public void testCopyOutFromEmptyJsonKeepsDefaults() { + assertEquals(PgCopyOutOptions.DEFAULT_AGGREGATION_THRESHOLD, + new PgCopyOutOptions(new JsonObject()).getAggregationThreshold()); + } + + @Test + public void testCopyOutRejectsInvalidThreshold() { + expectIllegalArgument(() -> new PgCopyOutOptions().setAggregationThreshold(0)); + expectIllegalArgument(() -> new PgCopyOutOptions().setAggregationThreshold(-1)); + assertEquals(1, new PgCopyOutOptions().setAggregationThreshold(1).getAggregationThreshold()); + } + + private static void expectIllegalArgument(Runnable action) { + try { + action.run(); + fail("expected an IllegalArgumentException"); + } catch (IllegalArgumentException expected) { + } + } + + private static void expectNullPointer(Runnable action) { + try { + action.run(); + fail("expected a NullPointerException"); + } catch (NullPointerException expected) { + } + } + + private static CopyOutStreamImpl copyOut() { + return new CopyOutStreamImpl(currentContext(), new TestReadStream()); + } + + private static CopyInStreamImpl copyIn() { + return new CopyInStreamImpl(currentContext(), new PgCopyInOptions()); + } + + + private static final class TestReadStream implements ReadStream { + + int pauseCalls; + int resumeCalls; + + @Override + public ReadStream exceptionHandler(Handler handler) { + return this; + } + + @Override + public ReadStream handler(Handler handler) { + return this; + } + + @Override + public ReadStream pause() { + pauseCalls++; + return this; + } + + @Override + public ReadStream resume() { + resumeCalls++; + return this; + } + + @Override + public ReadStream fetch(long amount) { + return this; + } + + @Override + public ReadStream endHandler(Handler endHandler) { + return this; + } + } +} diff --git a/vertx-pg-client/src/test/java/io/vertx/tests/pgclient/PgCopyTLSTest.java b/vertx-pg-client/src/test/java/io/vertx/tests/pgclient/PgCopyTLSTest.java new file mode 100644 index 000000000..a454b2c81 --- /dev/null +++ b/vertx-pg-client/src/test/java/io/vertx/tests/pgclient/PgCopyTLSTest.java @@ -0,0 +1,236 @@ +/* + * Copyright (c) 2011-2026 Contributors to the Eclipse Foundation + * + * This program and the accompanying materials are made available under the + * terms of the Eclipse Public License 2.0 which is available at + * http://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0 + * which is available at https://www.apache.org/licenses/LICENSE-2.0. + * + * SPDX-License-Identifier: EPL-2.0 OR Apache-2.0 + */ + +package io.vertx.tests.pgclient; + +import io.vertx.core.Future; +import io.vertx.core.Promise; +import io.vertx.core.Vertx; +import io.vertx.core.buffer.Buffer; +import io.vertx.core.net.ClientSSLOptions; +import io.vertx.ext.unit.TestContext; +import io.vertx.ext.unit.junit.VertxUnitRunner; +import io.vertx.pgclient.PgConnectOptions; +import io.vertx.pgclient.PgConnection; +import io.vertx.pgclient.PgCopyInOptions; +import io.vertx.pgclient.PgCopyOut; +import io.vertx.pgclient.PgCopyOutOptions; +import io.vertx.pgclient.SslMode; +import io.vertx.tests.pgclient.junit.ContainerPgRule; +import org.junit.After; +import org.junit.Before; +import org.junit.ClassRule; +import org.junit.Test; +import org.junit.runner.RunWith; + +import java.nio.charset.StandardCharsets; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Function; + +/** + * COPY over a TLS connection. + *

+ * This is worth covering on its own because the COPY IN sink drives backpressure from + * {@code channel().isWritable()} and installs its handler relative to the codec, while an + * {@code SslHandler} sits in the same pipeline doing its own buffering and re-framing. + */ +@RunWith(VertxUnitRunner.class) +public class PgCopyTLSTest { + + @ClassRule + public static ContainerPgRule rule = new ContainerPgRule().ssl(true); + + private Vertx vertx; + private PgConnectOptions options; + + @Before + public void setup() { + vertx = Vertx.vertx(); + options = new PgConnectOptions(rule.options()) + .setSslMode(SslMode.REQUIRE) + .setSslOptions(new ClientSSLOptions().setTrustAll(true)); + } + + @After + public void tearDown(TestContext ctx) { + vertx.close().onComplete(ctx.asyncAssertSuccess()); + } + + @Test + public void testCopyInOverTls(TestContext ctx) { + run(ctx, conn -> table(conn) + .compose(v -> conn.copyIn(COPY_IN)) + .compose(in -> in.write(buf("1,alpha\n2,beta\n")).compose(x -> in.end()).compose(x -> in.completion())) + .map(rows -> { + ctx.assertEquals(2, rows); + return null; + })); + } + + @Test + public void testCopyOutOverTls(TestContext ctx) { + run(ctx, conn -> conn.copyOut("COPY (SELECT generate_series(1, 5)) TO STDOUT") + .compose(this::drain) + .map(text -> { + ctx.assertEquals("1\n2\n3\n4\n5\n", text); + return null; + })); + } + + @Test + public void testLargeCopyInOverTlsCrossesTlsRecordBoundaries(TestContext ctx) { + // comfortably larger than the 16k TLS record size, so the payload spans many records + int rows = 50_000; + StringBuilder sb = new StringBuilder(); + for (int i = 1; i <= rows; i++) { + sb.append(i).append(",value-").append(i).append('\n'); + } + String payload = sb.toString(); + + run(ctx, conn -> table(conn) + .compose(v -> conn.copyIn(COPY_IN, new PgCopyInOptions())) + .compose(in -> in.write(buf(payload)).compose(x -> in.end()).compose(x -> in.completion())) + .map(count -> { + ctx.assertEquals(rows, count); + return null; + })); + } + + @Test + public void testLargeCopyOutOverTls(TestContext ctx) { + int rows = 50_000; + run(ctx, conn -> conn.copyOut("COPY (SELECT generate_series(1, " + rows + ")) TO STDOUT", + new PgCopyOutOptions().setAggregationThreshold(32 * 1024)) + .compose(this::drain) + .map(text -> { + StringBuilder expected = new StringBuilder(); + for (int i = 1; i <= rows; i++) { + expected.append(i).append('\n'); + } + ctx.assertEquals(expected.toString(), text); + return null; + })); + } + + @Test + public void testCopyInBackpressureOverTls(TestContext ctx) { + // a small write queue over TLS exercises the writability signalling through the SslHandler + run(ctx, conn -> table(conn) + .compose(v -> conn.copyIn(COPY_IN, new PgCopyInOptions() + .setChunkSize(1024))) + .compose(in -> { + in.setWriteQueueMaxSize(4096); + Future writes = Future.succeededFuture(); + for (int i = 1; i <= 5000; i++) { + int n = i; + writes = writes.compose(v -> in.write(buf(n + ",value-" + n + "\n"))); + } + return writes.compose(v -> in.end()).compose(v -> in.completion()); + }) + .map(rows -> { + ctx.assertEquals(5000, rows); + return null; + })); + } + + @Test + public void testCopyOutPausedFetchOverTls(TestContext ctx) { + run(ctx, conn -> conn.copyOut("COPY (SELECT generate_series(1, 200)) TO STDOUT") + .compose(out -> { + Promise promise = Promise.promise(); + StringBuilder sb = new StringBuilder(); + AtomicInteger seen = new AtomicInteger(); + out.pause(); + out.handler(b -> { + sb.append(b.toString(StandardCharsets.UTF_8)); + seen.incrementAndGet(); + out.fetch(1); + }); + out.endHandler(v -> promise.complete(sb.toString())); + out.exceptionHandler(promise::tryFail); + out.fetch(1); + return promise.future().compose(text -> out.completion().map(rows -> { + ctx.assertEquals(200, rows); + StringBuilder expected = new StringBuilder(); + for (int i = 1; i <= 200; i++) { + expected.append(i).append('\n'); + } + ctx.assertEquals(expected.toString(), text); + return (Void) null; + })); + })); + } + + @Test + public void testAbortOverTlsRollsBackAndKeepsConnectionUsable(TestContext ctx) { + run(ctx, conn -> table(conn) + .compose(v -> conn.copyIn(COPY_IN)) + .compose(in -> in.write(buf("1,alpha\n")) + .compose(x -> in.abort("tls abort")) + .compose(x -> in.completion().otherwise(-1))) + .compose(v -> conn.query("SELECT count(*) AS c FROM copy_tls").execute()) + .map(rs -> { + ctx.assertEquals(0, rs.iterator().next().getInteger("c")); + return null; + })); + } + + @Test + public void testCopyInRestoresWatermarksOverTls(TestContext ctx) { + run(ctx, conn -> table(conn) + .compose(v -> conn.copyIn(COPY_IN)) + .compose(in -> in.write(buf("1,alpha\n")).compose(x -> in.end()).compose(x -> in.completion())) + // a second COPY on the same TLS connection must still work after the sink detached + .compose(v -> conn.copyIn(COPY_IN)) + .compose(in -> in.write(buf("2,beta\n")).compose(x -> in.end()).compose(x -> in.completion())) + .compose(v -> conn.query("SELECT count(*) AS c FROM copy_tls").execute()) + .map(rs -> { + ctx.assertEquals(2, rs.iterator().next().getInteger("c")); + return null; + })); + } + + // ---------------------------------------------------------------- helpers + + private static final String COPY_IN = "COPY copy_tls (id, val) FROM STDIN WITH (FORMAT csv)"; + + private void run(TestContext ctx, Function> body) { + PgConnection.connect(vertx, options) + .compose(conn -> { + Future op; + try { + op = body.apply(conn); + } catch (Throwable t) { + op = Future.failedFuture(t); + } + return op.eventually(conn::close); + }) + .onComplete(ctx.asyncAssertSuccess()); + } + + private static Future table(PgConnection conn) { + return conn.query("CREATE TEMP TABLE copy_tls (id INT PRIMARY KEY, val TEXT NOT NULL)") + .execute().mapEmpty(); + } + + private Future drain(PgCopyOut out) { + Promise promise = Promise.promise(); + Buffer acc = Buffer.buffer(); + out.handler(acc::appendBuffer); + out.endHandler(v -> promise.tryComplete(acc)); + out.exceptionHandler(promise::tryFail); + return promise.future().map(b -> b.toString(StandardCharsets.UTF_8)); + } + + private static Buffer buf(String value) { + return Buffer.buffer(value, StandardCharsets.UTF_8.name()); + } +} diff --git a/vertx-pg-client/src/test/java/io/vertx/tests/pgclient/PgCopyTest.java b/vertx-pg-client/src/test/java/io/vertx/tests/pgclient/PgCopyTest.java new file mode 100644 index 000000000..04d0e7959 --- /dev/null +++ b/vertx-pg-client/src/test/java/io/vertx/tests/pgclient/PgCopyTest.java @@ -0,0 +1,1716 @@ +/* + * Copyright (c) 2011-2026 Contributors to the Eclipse Foundation + * + * This program and the accompanying materials are made available under the + * terms of the Eclipse Public License 2.0 which is available at + * http://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0 + * which is available at https://www.apache.org/licenses/LICENSE-2.0. + * + * SPDX-License-Identifier: EPL-2.0 OR Apache-2.0 + */ + +package io.vertx.tests.pgclient; + +import io.netty.channel.ChannelConfig; +import io.netty.channel.WriteBufferWaterMark; +import io.vertx.core.Future; +import io.vertx.core.Promise; +import io.vertx.core.Vertx; +import io.vertx.core.buffer.Buffer; +import io.vertx.ext.unit.Async; +import io.vertx.ext.unit.TestContext; +import io.vertx.pgclient.PgBuilder; +import io.vertx.pgclient.PgConnection; +import io.vertx.pgclient.PgCopyIn; +import io.vertx.pgclient.PgCopyInOptions; +import io.vertx.pgclient.PgCopyOut; +import io.vertx.pgclient.PgCopyOutOptions; +import io.vertx.pgclient.PgException; +import io.vertx.pgclient.impl.PgConnectionImpl; +import io.vertx.pgclient.impl.PgSocketConnection; +import io.vertx.sqlclient.Pool; +import io.vertx.sqlclient.PoolOptions; +import io.vertx.sqlclient.Row; +import io.vertx.sqlclient.RowSet; +import java.util.concurrent.atomic.AtomicBoolean; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.BiFunction; +import java.util.function.Function; + +public class PgCopyTest extends PgTestBase { + + private Vertx vertx; + + @Before + public void setup() throws Exception { + super.setup(); + vertx = Vertx.vertx(); + } + + @After + public void tearDown(TestContext ctx) { + vertx.close().onComplete(ctx.asyncAssertSuccess()); + } + + @Test + public void testCopyInDefaultDirectImmediateTransport(TestContext ctx) { + PgCopyInOptions options = new PgCopyInOptions(); + assertCopyInCsv(ctx, options, Arrays.asList("1:alpha", "2:beta", "3:gamma")); + } + + @Test + public void testCopyOutDefaultAggregationThreshold(TestContext ctx) { + PgCopyOutOptions options = new PgCopyOutOptions(); + ctx.assertEquals(1, options.getAggregationThreshold()); + } + + @Test + public void testCopyOptionsRejectInvalidValues(TestContext ctx) { + assertThrows(ctx, IllegalArgumentException.class, () -> new PgCopyInOptions().setChunkSize(0)); + assertThrows(ctx, IllegalArgumentException.class, () -> new PgCopyInOptions().setChunkSize(-1)); + assertThrows(ctx, IllegalArgumentException.class, () -> new PgCopyOutOptions().setAggregationThreshold(0)); + assertThrows(ctx, IllegalArgumentException.class, () -> new PgCopyOutOptions().setAggregationThreshold(-1)); + } + + @Test + public void testCopyInCoalescingTransportSmallChunks(TestContext ctx) { + assertCopyInCsv(ctx, new PgCopyInOptions() + .setChunkSize(8), Arrays.asList("1:alpha", "2:beta", "3:gamma")); + } + + @Test + public void testCopyInCoalescingTransportLargeBuffer(TestContext ctx) { + assertCopyInCsvSingleLargeBuffer(ctx, new PgCopyInOptions() + .setChunkSize(1024 * 1024)); + } + + @Test + public void testCopyInCoalescingHandoffSmallChunks(TestContext ctx) { + assertCopyInCsv(ctx, new PgCopyInOptions() + .setChunkSize(8), Arrays.asList("1:alpha", "2:beta", "3:gamma")); + } + + @Test + public void testCopyInCoalescingHandoffLargeBuffer(TestContext ctx) { + assertCopyInCsvSingleLargeBuffer(ctx, new PgCopyInOptions() + .setChunkSize(1024 * 1024)); + } + + @Test + public void testCopyInCoalescingSequentialComposedWrites(TestContext ctx) { + assertCopyInSequentialComposedWrites(ctx); + } + + @Test + public void testCopyInCoalescingEmptyWriteThenEnd(TestContext ctx) { + Async async = ctx.async(); + + withConnection(conn -> createTextTable(conn, "copy_test") + .compose(v -> conn.copyIn( + "COPY copy_test (id, val) FROM STDIN WITH (FORMAT csv)", + new PgCopyInOptions() + .setChunkSize(1024 * 1024))) + .compose(in -> in.write(Buffer.buffer()).compose(v -> in.end())) + .compose(v -> fetchCount(conn, "copy_test")) + .map(count -> { + ctx.assertEquals(0, count); + return null; + }) + ).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + @Test + public void testCopyInRestoresCustomChannelWatermarkAfterSuccess(TestContext ctx) { + Async async = ctx.async(); + + withConnection(conn -> createTextTable(conn, "copy_test").compose(v -> { + ChannelConfig config = channelConfig(conn); + WriteBufferWaterMark original = new WriteBufferWaterMark(123456, 234567); + config.setWriteBufferWaterMark(original); + + return conn.copyIn("COPY copy_test (id, val) FROM STDIN WITH (FORMAT csv)") + .compose(in -> { + assertWatermark(ctx, config.getWriteBufferWaterMark(), 4 * 1024 * 1024, 8 * 1024 * 1024); + in.setWriteQueueMaxSize(512 * 1024); + assertWatermark(ctx, config.getWriteBufferWaterMark(), 256 * 1024, 512 * 1024); + return in.write(buf("1,alpha\n")).compose(x -> in.end()); + }) + .map(x -> { + ctx.assertTrue(config.getWriteBufferWaterMark() == original); + return null; + }); + })).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + @Test + public void testCopyInRestoresCustomChannelWatermarkAfterAbort(TestContext ctx) { + Async async = ctx.async(); + + withConnection(conn -> createTextTable(conn, "copy_test").compose(v -> { + ChannelConfig config = channelConfig(conn); + WriteBufferWaterMark original = new WriteBufferWaterMark(123457, 234568); + config.setWriteBufferWaterMark(original); + + return conn.copyIn("COPY copy_test (id, val) FROM STDIN WITH (FORMAT csv)") + .compose(in -> in.abort("rollback") + .compose(x -> expectFailure(in.completion(), "abort should fail completion"))) + .compose(x -> conn.query("SELECT 1").execute()) + .map(x -> { + ctx.assertTrue(config.getWriteBufferWaterMark() == original); + return null; + }); + })).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + @Test + public void testCopyInRestoresCustomChannelWatermarkAfterServerError(TestContext ctx) { + Async async = ctx.async(); + + withConnection(conn -> createTextTable(conn, "copy_test").compose(v -> { + ChannelConfig config = channelConfig(conn); + WriteBufferWaterMark original = new WriteBufferWaterMark(123458, 234569); + config.setWriteBufferWaterMark(original); + + return conn.copyIn("COPY copy_test (id, val) FROM STDIN WITH (FORMAT csv)") + .compose(in -> in.write(buf("not-an-int,alpha\n")) + .compose(x -> expectFailure(in.end(), "server-side COPY error should fail end"))) + .compose(x -> conn.query("SELECT 1").execute()) + .map(x -> { + ctx.assertTrue(config.getWriteBufferWaterMark() == original); + return null; + }); + })).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + @Test + public void testCopyInDirectDelayedFlush(TestContext ctx) { + assertCopyInCsv(ctx, new PgCopyInOptions(), Arrays.asList("1:alpha", "2:beta", "3:gamma")); + } + + @Test + public void testCopyInDirectImmediateFlush(TestContext ctx) { + assertCopyInCsv(ctx, new PgCopyInOptions(), Arrays.asList("1:alpha", "2:beta", "3:gamma")); + } + + @Test + public void testCopyInFragmentedCsvLineAcrossManyWrites(TestContext ctx) { + Async async = ctx.async(); + + withConnection(conn -> createTextTable(conn, "copy_test") + .compose(v -> conn.copyIn("COPY copy_test (id, val) FROM STDIN WITH (FORMAT csv)", new PgCopyInOptions().setChunkSize(4))) + .compose(in -> { + List> writes = new ArrayList<>(); + writes.add(in.write(buf("1,"))); + writes.add(in.write(buf("al"))); + writes.add(in.write(buf("pha\n2"))); + writes.add(in.write(buf(",be"))); + writes.add(in.write(buf("ta\n"))); + writes.add(in.write(buf("3,gamma\n"))); + + Future completion = in.end(); + return Future.all(writes).compose(v -> completion); + }) + .compose(v -> fetchTextRows(conn, "copy_test")) + .map(rows -> { + ctx.assertEquals(Arrays.asList("1:alpha", "2:beta", "3:gamma"), rows); + return null; + }) + ).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + @Test + public void testCopyInCsvHeaderAndEscaping(TestContext ctx) { + Async async = ctx.async(); + + withConnection(conn -> createTextTable(conn, "copy_test") + .compose(v -> conn.copyIn("COPY copy_test (id, val) FROM STDIN WITH (FORMAT csv, HEADER true)")) + .compose(in -> { + Future write = in.write(buf( + "id,val\n" + + "1,alpha\n" + + "2,\"hello, world\"\n" + + "3,\"he said \"\"yo\"\"\"\n")); + Future completion = in.end(); + return write.compose(v -> completion); + }) + .compose(v -> fetchTextRows(conn, "copy_test")) + .map(rows -> { + ctx.assertEquals(Arrays.asList("1:alpha", "2:hello, world", "3:he said \"yo\""), rows); + return null; + }) + ).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + @Test + public void testCopyInTextFormat(TestContext ctx) { + Async async = ctx.async(); + + withConnection(conn -> createTextTable(conn, "copy_test") + .compose(v -> conn.copyIn("COPY copy_test (id, val) FROM STDIN WITH (FORMAT text)")) + .compose(in -> { + Future write = in.write(buf("1\talpha\n2\tbeta\n3\tgamma\n")); + Future completion = in.end(); + return write.compose(v -> completion); + }) + .compose(v -> fetchTextRows(conn, "copy_test")) + .map(rows -> { + ctx.assertEquals(Arrays.asList("1:alpha", "2:beta", "3:gamma"), rows); + return null; + }) + ).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + @Test + public void testCopyInEmptyEnd(TestContext ctx) { + Async async = ctx.async(); + + withConnection(conn -> createTextTable(conn, "copy_test") + .compose(v -> conn.copyIn("COPY copy_test (id, val) FROM STDIN WITH (FORMAT csv)")) + .compose(PgCopyIn::end) + .compose(v -> fetchCount(conn, "copy_test")) + .map(count -> { + ctx.assertEquals(0, count); + return null; + }) + ).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + @Test + public void testCopyInWriteAfterEndFailsAndDoesNotInsert(TestContext ctx) { + Async async = ctx.async(); + + withConnection(conn -> createTextTable(conn, "copy_test") + .compose(v -> conn.copyIn("COPY copy_test (id, val) FROM STDIN WITH (FORMAT csv)")) + .compose(in -> { + Future completion = in.end(); + return expectFailure(in.write(buf("1,late\n")), "write after end should fail") + .compose(v -> completion); + }) + .compose(v -> fetchCount(conn, "copy_test")) + .map(count -> { + ctx.assertEquals(0, count); + return null; + }) + ).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + @Test + public void testCopyInAbortAfterEndFailsAndDoesNotSendCopyFail(TestContext ctx) { + Async async = ctx.async(); + + withConnection(conn -> createTextTable(conn, "copy_test") + .compose(v -> conn.copyIn("COPY copy_test (id, val) FROM STDIN WITH (FORMAT csv)")) + .compose(in -> { + Future completion = in.end(); + return expectFailure(in.abort("too late"), "abort after end should fail") + .compose(v -> completion); + }) + .compose(v -> conn.query("SELECT 1").execute().mapEmpty()) + ).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + @Test + public void testCopyInWriteAfterAbortFailsAndConnectionIsReusable(TestContext ctx) { + Async async = ctx.async(); + + withConnection(conn -> createTextTable(conn, "copy_test") + .compose(v -> conn.copyIn("COPY copy_test (id, val) FROM STDIN WITH (FORMAT csv)")) + .compose(in -> in.abort("stop") + .compose(v -> expectFailure(in.write(buf("1,late\n")), "write after abort should fail")) + .compose(v -> expectFailure(in.completion(), "abort should fail completion"))) + .compose(v -> conn.query("SELECT 1").execute().mapEmpty()) + ).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + @Test + public void testCopyInAbortAfterBufferedCoalescingDataRollsBackCopy(TestContext ctx) { + Async async = ctx.async(); + + withConnection(conn -> createTextTable(conn, "copy_test") + .compose(v -> conn.copyIn( + "COPY copy_test (id, val) FROM STDIN WITH (FORMAT csv)", + new PgCopyInOptions() + .setChunkSize(1024 * 1024))) + .compose(in -> { + Future write = in.write(buf("1,alpha\n2,beta\n")); + Future abort = in.abort("rollback copy"); + + return expectFailure(write, "buffered coalescing write should fail after abort") + .compose(v -> abort) + .compose(v -> expectFailure(in.completion(), "abort should fail completion")); + }) + .compose(v -> fetchCount(conn, "copy_test")) + .map(count -> { + ctx.assertEquals(0, count); + return null; + }) + ).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + @Test + public void testCopyInServerErrorFailsCompletionAndConnectionIsReusable(TestContext ctx) { + Async async = ctx.async(); + + withConnection(conn -> createTextTable(conn, "copy_test") + .compose(v -> conn.copyIn("COPY copy_test (id, val) FROM STDIN WITH (FORMAT csv)")) + .compose(in -> { + Future write = in.write(buf("not-an-int,alpha\n")); + Future completion = in.end(); + + return write.otherwiseEmpty() + .compose(v -> expectFailure(completion, "server-side COPY error should fail completion")); + }) + .compose(v -> conn.query("SELECT 1").execute().mapEmpty()) + ).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + @Test + public void testCopyInFollowingQueryWaitsForCopyCompletion(TestContext ctx) { + Async async = ctx.async(); + + withConnection(conn -> createTextTable(conn, "copy_test") + .compose(v -> conn.copyIn("COPY copy_test (id, val) FROM STDIN WITH (FORMAT csv)", new PgCopyInOptions().setChunkSize(1024 * 1024))) + .compose(in -> { + Future> select = conn.query("SELECT count(*) AS c FROM copy_test").execute(); + ctx.assertFalse(select.isComplete()); + + Future write = in.write(buf("1,alpha\n2,beta\n")); + Future completion = in.end(); + + return Future.all(write, completion).compose(v -> select); + }) + .map(rows -> { + ctx.assertEquals(2, rows.iterator().next().getInteger("c")); + return null; + }) + ).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + @Test + public void testCopyInRejectsInvalidSql(TestContext ctx) { + Async async = ctx.async(); + + withConnection(conn -> Future.all( + // rejected as a failed future, never thrown at the caller + assertFailsWith(ctx, NullPointerException.class, conn.copyIn(null)), + assertFailsWith(ctx, IllegalArgumentException.class, conn.copyIn("SELECT 1")), + assertFailsWith(ctx, IllegalArgumentException.class, conn.copyIn("COPY (SELECT 1) TO STDOUT"))) + .compose(v -> conn.query("SELECT 1").execute().mapEmpty()) + ).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + @Test + public void testCopyOutDefaultUnpaused(TestContext ctx) { + Async async = ctx.async(); + + PgCopyOutOptions options = new PgCopyOutOptions(); + ctx.assertEquals(1, options.getAggregationThreshold()); + + withConnection(conn -> createTextTable(conn, "copy_test") + .compose(v -> insertTextRows(conn, "copy_test")) + .compose(v -> conn.copyOut("COPY (SELECT id, val FROM copy_test ORDER BY id) TO STDOUT WITH (FORMAT csv)")) + .compose(out -> collectFlowing(out) + .compose(csv -> out.completion().map(rowCount -> { + ctx.assertEquals("1,alpha\n2,beta\n3,gamma\n", normalize(csv)); + ctx.assertEquals(3, rowCount); + return null; + }))) + ).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + @Test + public void testCopyOutExplicitAggregationThresholdOneCsvHeader(TestContext ctx) { + Async async = ctx.async(); + + withConnection(conn -> createTextTable(conn, "copy_test") + .compose(v -> insertTextRows(conn, "copy_test")) + .compose(v -> conn.copyOut( + "COPY (SELECT id, val FROM copy_test ORDER BY id) TO STDOUT WITH (FORMAT csv, HEADER true)", + new PgCopyOutOptions().setAggregationThreshold(1))) + .compose(out -> collect(out).map(csv -> { + ctx.assertEquals("id,val\n1,alpha\n2,beta\n3,gamma\n", normalize(csv)); + return null; + })) + ).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + @Test + public void testCopyOutLargeAggregationThresholdUnpaused(TestContext ctx) { + Async async = ctx.async(); + + withConnection(conn -> createTextTable(conn, "copy_test") + .compose(v -> insertTextRows(conn, "copy_test")) + .compose(v -> conn.copyOut( + "COPY (SELECT id, val FROM copy_test ORDER BY id) TO STDOUT WITH (FORMAT csv)", + new PgCopyOutOptions().setAggregationThreshold(1024 * 1024))) + .compose(out -> collectFlowing(out).map(csv -> { + ctx.assertEquals("1,alpha\n2,beta\n3,gamma\n", normalize(csv)); + return null; + })) + ).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + @Test + public void testCopyOutTextFormat(TestContext ctx) { + Async async = ctx.async(); + + withConnection(conn -> createTextTable(conn, "copy_test") + .compose(v -> insertTextRows(conn, "copy_test")) + .compose(v -> conn.copyOut("COPY (SELECT id, val FROM copy_test ORDER BY id) TO STDOUT WITH (FORMAT text)")) + .compose(out -> collect(out).map(text -> { + ctx.assertEquals("1\talpha\n2\tbeta\n3\tgamma\n", normalize(text)); + return null; + })) + ).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + @Test + public void testCopyOutPausedFetchOneByOne(TestContext ctx) { + Async async = ctx.async(); + + PgCopyOutOptions options = new PgCopyOutOptions(); + ctx.assertEquals(1, options.getAggregationThreshold()); + + withConnection(conn -> createTextTable(conn, "copy_test") + .compose(v -> insertTextRows(conn, "copy_test")) + .compose(v -> conn.copyOut( + "COPY (SELECT id, val FROM copy_test ORDER BY id) TO STDOUT WITH (FORMAT csv)")) + .compose(out -> { + StringBuilder sb = new StringBuilder(); + Promise drained = Promise.promise(); + + out.pause(); + out.handler(buffer -> { + sb.append(buffer.toString(StandardCharsets.UTF_8)); + out.fetch(1); + }); + out.exceptionHandler(drained::tryFail); + out.endHandler(v -> drained.tryComplete(sb.toString())); + out.completion().onFailure(drained::tryFail); + out.fetch(1); + + return drained.future(); + }) + .map(csv -> { + ctx.assertEquals("1,alpha\n2,beta\n3,gamma\n", normalize(csv)); + return null; + }) + ).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + @Test + public void testCopyOutFetchOneByOneAcrossTransportWatermarks(TestContext ctx) { + Async async = ctx.async(); + + withConnection(conn -> conn.copyOut( + "COPY (SELECT generate_series(1, 256)) TO STDOUT WITH (FORMAT text)") + .compose(out -> { + AtomicInteger chunks = new AtomicInteger(); + Promise drained = Promise.promise(); + + out.handler(buffer -> { + chunks.incrementAndGet(); + out.fetch(1); + }); + out.exceptionHandler(drained::tryFail); + out.endHandler(v -> drained.tryComplete(chunks.get())); + out.completion().onFailure(drained::tryFail); + out.fetch(1); + + return drained.future().compose(count -> out.completion().map(rowCount -> { + ctx.assertEquals(256, count); + ctx.assertEquals(256, rowCount); + return null; + })); + }) + ).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + @Test + public void testCopyOutServerErrorFailsReadyFutureAndConnectionIsReusable(TestContext ctx) { + Async async = ctx.async(); + + withConnection(conn -> conn.copyOut("COPY (SELECT * FROM missing_copy_table) TO STDOUT WITH (FORMAT csv)") + .compose(v -> Future.failedFuture("COPY OUT should have failed"), err -> Future.succeededFuture()) + .compose(v -> conn.query("SELECT 1").execute().mapEmpty()) + ).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + @Test + public void testCopyOutRejectsInvalidSql(TestContext ctx) { + Async async = ctx.async(); + + withConnection(conn -> createTextTable(conn, "copy_test") + .compose(v -> Future.all( + // rejected as a failed future, never thrown at the caller + assertFailsWith(ctx, NullPointerException.class, conn.copyOut(null)), + assertFailsWith(ctx, IllegalArgumentException.class, conn.copyOut("SELECT 1")), + assertFailsWith(ctx, IllegalArgumentException.class, conn.copyOut("COPY copy_test FROM STDIN")))) + .compose(v -> conn.query("SELECT 1").execute().mapEmpty()) + ).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + @Test + public void testCopyOutUnexpectedResponseFromPlainQueryClosesConnection(TestContext ctx) { + assertUnexpectedCopyQueryClosesConnection(ctx, "COPY (SELECT 1) TO STDOUT WITH (FORMAT csv)"); + } + + @Test + public void testCopyInUnexpectedResponseFromPlainQueryClosesConnection(TestContext ctx) { + Async async = ctx.async(); + + withConnection(conn -> createTextTable(conn, "copy_test") + .compose(v -> unexpectedCopyQuery(conn, "COPY copy_test (id, val) FROM STDIN WITH (FORMAT csv)")) + ).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + @Test + public void testCopyInAbortOnPooledConnectionRecycle(TestContext ctx) { + Async async = ctx.async(); + Pool pool = PgBuilder.pool(b -> b.connectingTo(options).with(new PoolOptions().setMaxSize(1)).using(vertx)); + + pool.getConnection() + .compose(conn -> { + PgConnection pg = (PgConnection) conn; + return createTextTable(pg, "copy_test") + .compose(v -> pg.copyIn("COPY copy_test (id, val) FROM STDIN WITH (FORMAT csv)")) + .compose(copy -> copy.write(buf("1,alpha\n")).compose(v -> pg.close())); + }) + .compose(v -> pool.getConnection()) + .compose(conn -> { + PgConnection pg = (PgConnection) conn; + return fetchCount(pg, "copy_test") + .map(count -> { + ctx.assertEquals(0, count); + return null; + }) + .eventually(pg::close); + }) + .eventually(pool::close) + .onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + @Test + public void testPendingCopyInAbortOnPooledConnectionRecycle(TestContext ctx) { + Async async = ctx.async(); + Pool pool = PgBuilder.pool(b -> b.connectingTo(options).with(new PoolOptions().setMaxSize(1)).using(vertx)); + + pool.getConnection() + .compose(conn -> { + PgConnection pg = (PgConnection) conn; + return createTextTable(pg, "copy_test") + .compose(v -> { + pg.copyIn("COPY copy_test (id, val) FROM STDIN WITH (FORMAT csv)"); + return pg.close(); + }); + }) + .compose(v -> pool.getConnection()) + .compose(conn -> { + PgConnection pg = (PgConnection) conn; + return fetchCount(pg, "copy_test") + .map(count -> { + ctx.assertEquals(0, count); + return null; + }) + .eventually(pg::close); + }) + .eventually(pool::close) + .onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + @Test + public void testCopyOutCompletesBeforePooledConnectionRecycle(TestContext ctx) { + Async async = ctx.async(); + Pool pool = PgBuilder.pool(b -> b.connectingTo(options).with(new PoolOptions().setMaxSize(1)).using(vertx)); + + pool.getConnection() + .compose(conn -> { + PgConnection pg = (PgConnection) conn; + return pg.copyOut("COPY (SELECT generate_series(1, 5000)) TO STDOUT WITH (FORMAT csv)") + .compose(copy -> pg.close()); + }) + .compose(v -> pool.getConnection()) + .compose(conn -> { + PgConnection pg = (PgConnection) conn; + return pg.query("SELECT 1").execute().mapEmpty().eventually(pg::close); + }) + .eventually(pool::close) + .onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + @Test + public void testCopyOutPipeToCopyInCsvAcrossConnections(TestContext ctx) { + Async async = ctx.async(); + + withTwoConnections((src, dst) -> createTextTable(src, "copy_src") + .compose(v -> createTextTable(dst, "copy_dst")) + .compose(v -> insertTextRows(src, "copy_src")) + .compose(v -> src.copyOut("COPY (SELECT id, val FROM copy_src ORDER BY id) TO STDOUT WITH (FORMAT csv)")) + .compose(out -> dst.copyIn("COPY copy_dst (id, val) FROM STDIN WITH (FORMAT csv)") + .compose(in -> out.pipeTo(in) + .compose(v -> out.completion()) + .compose(v -> in.completion()))) + .compose(v -> fetchTextRows(dst, "copy_dst")) + .map(rows -> { + ctx.assertEquals(Arrays.asList("1:alpha", "2:beta", "3:gamma"), rows); + return null; + }) + ).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + @Test + public void testCopyOutPipeToCopyInBinaryAcrossConnections(TestContext ctx) { + Async async = ctx.async(); + + withTwoConnections((src, dst) -> createTextTable(src, "copy_src") + .compose(v -> createTextTable(dst, "copy_dst")) + .compose(v -> insertTextRows(src, "copy_src")) + .compose(v -> src.copyOut("COPY (SELECT id, val FROM copy_src ORDER BY id) TO STDOUT WITH (FORMAT binary)")) + .compose(out -> dst.copyIn("COPY copy_dst (id, val) FROM STDIN WITH (FORMAT binary)") + .compose(in -> out.pipeTo(in) + .compose(v -> out.completion()) + .compose(v -> in.completion()))) + .compose(v -> fetchTextRows(dst, "copy_dst")) + .map(rows -> { + ctx.assertEquals(Arrays.asList("1:alpha", "2:beta", "3:gamma"), rows); + return null; + }) + ).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + + // ================================================================ + // Contract: how the futures resolve, how the connection is guarded, + // how bad input is reported. + // ================================================================ + + // ---------------------------------------------------------------- completion semantics + + @Test + public void testCopyInCompletionMeansConnectionIsReady(TestContext ctx) { + run(ctx, conn -> table(conn) + .compose(v -> conn.copyIn(copyIn())) + .compose(in -> in.write(buf("1,alpha\n")).compose(x -> in.end()).map(x -> in)) + .compose(in -> in.completion()) + .compose(rows -> { + ctx.assertEquals(1, rows); + // completion() resolves at ReadyForQuery, so the connection takes a new command right away + return conn.query("SELECT 1 AS n").execute(); + }) + .map(rs -> { + ctx.assertEquals(1, rs.iterator().next().getInteger("n")); + return null; + })); + } + + @Test + public void testCopyOutCompletionMeansConnectionIsReady(TestContext ctx) { + run(ctx, conn -> conn.copyOut("COPY (SELECT generate_series(1, 3)) TO STDOUT") + .compose(out -> drain(out).compose(text -> { + ctx.assertEquals("1\n2\n3\n", text); + return out.completion(); + })) + .compose(rows -> { + ctx.assertEquals(3, rows); + return conn.query("SELECT 1 AS n").execute(); + }) + .map(rs -> { + ctx.assertEquals(1, rs.iterator().next().getInteger("n")); + return null; + })); + } + + @Test + public void testCopyInEndFutureResolvesWhenCopyCompletes(TestContext ctx) { + run(ctx, conn -> table(conn) + .compose(v -> conn.copyIn(copyIn())) + .compose(in -> in.write(buf("1,alpha\n2,beta\n")).compose(x -> in.end())) + .compose(v -> rows(conn)) + .map(rows -> { + ctx.assertEquals(2, rows.size()); + return null; + })); + } + + // ---------------------------------------------------------------- abort + + @Test + public void testAbortSurfacesTheServerErrorInCompletion(TestContext ctx) { + AtomicReference cause = new AtomicReference<>(); + run(ctx, conn -> table(conn) + .compose(v -> conn.copyIn(copyIn())) + .compose(in -> in.write(buf("1,alpha\n")) + .compose(x -> in.abort("stop right there")) + .compose(x -> in.completion().otherwise(err -> { + cause.set(err); + return -1; + }))) + .map(v -> { + Throwable err = cause.get(); + ctx.assertNotNull(err, "completion() should have failed"); + // the outcome is reported by PostgreSQL, not synthesised by the client + ctx.assertTrue(err instanceof PgException, "expected a PgException but got " + err); + ctx.assertTrue(err.getMessage().contains("stop right there"), + "server error should quote the abort message, was: " + err.getMessage()); + return null; + })); + } + + @Test + public void testAbortRollsBackEverythingWritten(TestContext ctx) { + run(ctx, conn -> table(conn) + .compose(v -> conn.copyIn(copyIn())) + .compose(in -> in.write(buf("1,alpha\n2,beta\n")) + .compose(x -> in.abort("rollback")) + .compose(x -> in.completion().otherwise(-1))) + .compose(v -> count(conn)) + .map(count -> { + ctx.assertEquals(0, count); + return null; + })); + } + + @Test + public void testAbortItselfSucceeds(TestContext ctx) { + run(ctx, conn -> table(conn) + .compose(v -> conn.copyIn(copyIn())) + .compose(in -> in.abort("done") + .compose(x -> in.completion().otherwise(-1))) + .mapEmpty()); + } + + // ---------------------------------------------------------------- one COPY per connection + + @Test + public void testSecondCopyInWhileFirstIsActiveIsRejected(TestContext ctx) { + run(ctx, conn -> table(conn) + .compose(v -> conn.copyIn(copyIn())) + .compose(in -> conn.copyIn(copyIn()) + .compose(second -> Future.failedFuture("a second COPY IN should have been rejected"), + err -> { + ctx.assertTrue(err instanceof IllegalStateException, "was " + err); + return Future.succeededFuture(); + }) + .compose(x -> in.abort("cleanup")) + .compose(x -> in.completion().otherwise(-1))) + .mapEmpty()); + } + + @Test + public void testCopyOutWhileCopyInIsActiveIsRejected(TestContext ctx) { + run(ctx, conn -> table(conn) + .compose(v -> conn.copyIn(copyIn())) + .compose(in -> conn.copyOut("COPY (SELECT 1) TO STDOUT") + .compose(out -> Future.failedFuture("COPY OUT should have been rejected"), + err -> { + ctx.assertTrue(err instanceof IllegalStateException, "was " + err); + return Future.succeededFuture(); + }) + .compose(x -> in.abort("cleanup")) + .compose(x -> in.completion().otherwise(-1))) + .mapEmpty()); + } + + @Test + public void testCopyInAfterPreviousOneCompletedSucceeds(TestContext ctx) { + run(ctx, conn -> table(conn) + .compose(v -> copyRow(conn, "1,alpha\n")) + .compose(v -> copyRow(conn, "2,beta\n")) + .compose(v -> copyRow(conn, "3,gamma\n")) + .compose(v -> rows(conn)) + .map(rows -> { + ctx.assertEquals(3, rows.size()); + return null; + })); + } + + @Test + public void testCopyOutAfterCopyInOnTheSameConnection(TestContext ctx) { + run(ctx, conn -> table(conn) + .compose(v -> copyRow(conn, "1,alpha\n")) + .compose(v -> conn.copyOut("COPY copy_contract TO STDOUT WITH (FORMAT csv)")) + .compose(this::drain) + .map(text -> { + ctx.assertEquals("1,alpha\n", text); + return null; + })); + } + + // ---------------------------------------------------------------- rejected SQL + + @Test + public void testInvalidSqlFailsTheFutureInsteadOfThrowing(TestContext ctx) { + run(ctx, conn -> { + // must not throw at call time + Future in = conn.copyIn("SELECT 1"); + return assertFailsWithType(ctx, in, IllegalArgumentException.class) + .compose(v -> assertFailsWithType(ctx, conn.copyOut("SELECT 1"), IllegalArgumentException.class)); + }); + } + + @Test + public void testNullSqlFailsTheFuture(TestContext ctx) { + run(ctx, conn -> Future.all( + assertFailsWithType(ctx, conn.copyIn(null), NullPointerException.class), + assertFailsWithType(ctx, conn.copyOut(null), NullPointerException.class) + ).mapEmpty()); + } + + @Test + public void testWrongDirectionIsRejected(TestContext ctx) { + run(ctx, conn -> table(conn) + .compose(v -> assertFailsWithType(ctx, conn.copyIn("COPY copy_contract TO STDOUT"), IllegalArgumentException.class)) + .compose(v -> assertFailsWithType(ctx, conn.copyOut("COPY copy_contract FROM STDIN"), IllegalArgumentException.class)) + // the connection survives both mismatches + .compose(v -> copyRow(conn, "1,alpha\n")) + .compose(v -> count(conn)) + .map(count -> { + ctx.assertEquals(1, count); + return null; + })); + } + + @Test + public void testMultipleStatementsAreRejected(TestContext ctx) { + run(ctx, conn -> assertFailsWithType(ctx, + conn.copyIn("COPY copy_contract FROM STDIN; DROP TABLE copy_contract"), + IllegalArgumentException.class).mapEmpty()); + } + + @Test + public void testSemicolonInsideAnOptionIsNotAStatementSeparator(TestContext ctx) { + run(ctx, conn -> table(conn) + .compose(v -> conn.copyIn("COPY copy_contract (id, val) FROM STDIN WITH (FORMAT csv, DELIMITER ';')")) + .compose(in -> in.write(buf("1;alpha\n2;beta\n")).compose(v -> in.end())) + .compose(v -> rows(conn)) + .map(rows -> { + ctx.assertEquals(Arrays.asList("1:alpha", "2:beta"), rows); + return null; + })); + } + + @Test + public void testSemicolonInsideAQuotedIdentifierIsNotAStatementSeparator(TestContext ctx) { + run(ctx, conn -> exec(conn, "CREATE TEMP TABLE \"odd;name\" (id INT)") + .compose(v -> conn.copyIn("COPY \"odd;name\" FROM STDIN WITH (FORMAT csv)")) + .compose(in -> in.write(buf("1\n")).compose(v -> in.end())) + .mapEmpty()); + } + + @Test + public void testATrailingStatementAfterAQuotedSemicolonIsStillRejected(TestContext ctx) { + run(ctx, conn -> assertFailsWithType(ctx, + conn.copyIn("COPY copy_contract FROM STDIN (DELIMITER ';'); DROP TABLE copy_contract"), + IllegalArgumentException.class).mapEmpty()); + } + + @Test + public void testKeywordsMayBeWrappedOverSeveralLines(TestContext ctx) { + run(ctx, conn -> table(conn) + .compose(v -> conn.copyIn("COPY copy_contract (id, val)\n FROM\n STDIN\n WITH (FORMAT csv)")) + .compose(in -> in.write(buf("1,alpha\n")).compose(v -> in.end())) + .compose(v -> count(conn)) + .map(count -> { + ctx.assertEquals(1, count); + return null; + })); + } + + @Test + public void testCopyOutKeywordsMayBeWrappedOverSeveralLines(TestContext ctx) { + run(ctx, conn -> table(conn) + .compose(v -> copyRow(conn, "1,alpha\n")) + .compose(v -> conn.copyOut("COPY copy_contract\n TO\n STDOUT\n WITH (FORMAT csv)")) + .compose(out -> drain(out).compose(data -> out.completion().map(rowCount -> { + ctx.assertEquals(1, rowCount); + ctx.assertEquals("1,alpha\n", data); + return null; + })))); + } + + @Test + public void testRejectedSqlLeavesTheConnectionUsable(TestContext ctx) { + run(ctx, conn -> table(conn) + .compose(v -> assertFailsWithType(ctx, conn.copyIn("SELECT 1"), IllegalArgumentException.class)) + // the guard must not have been taken, a real COPY still works + .compose(v -> copyRow(conn, "1,alpha\n")) + .compose(v -> count(conn)) + .map(count -> { + ctx.assertEquals(1, count); + return null; + })); + } + + // ---------------------------------------------------------------- data fidelity + + @Test + public void testRoundTripPreservesSpecialCharacters(TestContext ctx) { + String tricky = "a \"quoted\" value, with comma\nand newline"; + String csv = "1,\"" + tricky.replace("\"", "\"\"") + "\"\n"; + run(ctx, conn -> table(conn) + .compose(v -> copyRow(conn, csv)) + .compose(v -> conn.query("SELECT val FROM copy_contract WHERE id = 1").execute()) + .map(rs -> { + ctx.assertEquals(tricky, rs.iterator().next().getString("val")); + return null; + })); + } + + @Test + public void testRoundTripPreservesUnicode(TestContext ctx) { + String unicode = "héllo — Ω 日本語 🐘"; + run(ctx, conn -> table(conn) + .compose(v -> copyRow(conn, "1,\"" + unicode + "\"\n")) + .compose(v -> conn.query("SELECT val FROM copy_contract WHERE id = 1").execute()) + .map(rs -> { + ctx.assertEquals(unicode, rs.iterator().next().getString("val")); + return null; + })); + } + + @Test + public void testTextFormatNullMarker(TestContext ctx) { + run(ctx, conn -> exec(conn, "CREATE TEMP TABLE copy_nullable (id INT, val TEXT)") + .compose(v -> conn.copyIn("COPY copy_nullable FROM STDIN")) + .compose(in -> in.write(buf("1\t\\N\n")).compose(x -> in.end())) + .compose(v -> conn.query("SELECT val FROM copy_nullable WHERE id = 1").execute()) + .map(rs -> { + ctx.assertNull(rs.iterator().next().getString("val")); + return null; + })); + } + + @Test + public void testCopyOutOfAnEmptyTable(TestContext ctx) { + AtomicBoolean ended = new AtomicBoolean(); + run(ctx, conn -> table(conn) + .compose(v -> conn.copyOut("COPY copy_contract TO STDOUT WITH (FORMAT csv)")) + .compose(out -> drain(out, ended).compose(text -> { + ctx.assertEquals("", text); + return out.completion(); + })) + .map(rows -> { + ctx.assertEquals(0, rows); + ctx.assertTrue(ended.get(), "endHandler should have fired"); + return null; + })); + } + + @Test + public void testCopyInWithNoRows(TestContext ctx) { + run(ctx, conn -> table(conn) + .compose(v -> conn.copyIn(copyIn())) + .compose(in -> in.end().compose(x -> in.completion())) + .map(rows -> { + ctx.assertEquals(0, rows); + return null; + })); + } + + @Test + public void testLargeCopyInAndOutRoundTrip(TestContext ctx) { + int count = 20_000; + StringBuilder sb = new StringBuilder(); + for (int i = 1; i <= count; i++) { + sb.append(i).append(",value-").append(i).append('\n'); + } + String payload = sb.toString(); + + run(ctx, conn -> table(conn) + .compose(v -> conn.copyIn(copyIn(), new PgCopyInOptions())) + .compose(in -> in.write(buf(payload)).compose(x -> in.end()).compose(x -> in.completion())) + .compose(rows -> { + ctx.assertEquals(count, rows); + return conn.copyOut("COPY copy_contract TO STDOUT WITH (FORMAT csv)"); + }) + .compose(this::drain) + .map(text -> { + ctx.assertEquals(payload.length(), text.length()); + ctx.assertEquals(payload, text); + return null; + })); + } + + @Test + public void testCopyOutQueryForm(TestContext ctx) { + run(ctx, conn -> conn.copyOut("COPY (SELECT i, 'v' || i FROM generate_series(1, 5) AS i) TO STDOUT WITH (FORMAT csv)") + .compose(this::drain) + .map(text -> { + ctx.assertEquals("1,v1\n2,v2\n3,v3\n4,v4\n5,v5\n", text); + return null; + })); + } + + @Test + public void testBinaryFormatRoundTripOnOneConnection(TestContext ctx) { + run(ctx, conn -> table(conn) + .compose(v -> copyRow(conn, "1,alpha\n2,beta\n")) + .compose(v -> conn.copyOut("COPY copy_contract TO STDOUT WITH (FORMAT binary)")) + .compose(out -> collectBytes(out)) + .compose(binary -> exec(conn, "DELETE FROM copy_contract") + .compose(v -> conn.copyIn("COPY copy_contract FROM STDIN WITH (FORMAT binary)")) + .compose(in -> in.write(binary).compose(x -> in.end()).compose(x -> in.completion()))) + .compose(rows -> { + ctx.assertEquals(2, rows); + return rows(conn); + }) + .map(all -> { + ctx.assertEquals(2, all.size()); + ctx.assertEquals("1:alpha", all.get(0)); + return null; + })); + } + + // ---------------------------------------------------------------- stream surface + + @Test + public void testEndWithBufferWritesTheData(TestContext ctx) { + run(ctx, conn -> table(conn) + .compose(v -> conn.copyIn(copyIn())) + .compose(in -> in.end(buf("1,alpha\n")).compose(x -> in.completion())) + .map(rows -> { + ctx.assertEquals(1, rows); + return null; + })); + } + + @Test + public void testWriteQueueFullAndDrainHandler(TestContext ctx) { + AtomicInteger drains = new AtomicInteger(); + run(ctx, conn -> table(conn) + .compose(v -> conn.copyIn(copyIn(), new PgCopyInOptions() + .setChunkSize(64))) + .compose(in -> { + in.setWriteQueueMaxSize(128); + in.drainHandler(x -> drains.incrementAndGet()); + + List> writes = new ArrayList<>(); + for (int i = 1; i <= 2000; i++) { + writes.add(in.write(buf(i + ",value-" + i + "\n"))); + } + return Future.all(new ArrayList<>(writes)) + .compose(x -> in.end()) + .compose(x -> in.completion()); + }) + .map(rows -> { + ctx.assertEquals(2000, rows); + return null; + })); + } + + @Test + public void testSetWriteQueueMaxSizeIsClampedToAtLeastOne(TestContext ctx) { + run(ctx, conn -> table(conn) + .compose(v -> conn.copyIn(copyIn())) + .compose(in -> { + in.setWriteQueueMaxSize(0); + in.setWriteQueueMaxSize(-100); + return in.write(buf("1,alpha\n")).compose(x -> in.end()).compose(x -> in.completion()); + }) + .map(rows -> { + ctx.assertEquals(1, rows); + return null; + })); + } + + @Test + public void testCopyOutPauseFetchResume(TestContext ctx) { + run(ctx, conn -> conn.copyOut("COPY (SELECT generate_series(1, 50)) TO STDOUT") + .compose(out -> { + Promise promise = Promise.promise(); + StringBuilder sb = new StringBuilder(); + AtomicInteger chunks = new AtomicInteger(); + + out.pause(); + out.handler(b -> { + sb.append(b.toString(StandardCharsets.UTF_8)); + if (chunks.incrementAndGet() < 10) { + out.fetch(1); + } else { + out.resume(); + } + }); + out.endHandler(v -> promise.complete(sb.toString())); + out.exceptionHandler(promise::tryFail); + out.fetch(1); + + return promise.future().compose(text -> out.completion().map(rows -> { + ctx.assertEquals(50, rows); + StringBuilder expected = new StringBuilder(); + for (int i = 1; i <= 50; i++) { + expected.append(i).append('\n'); + } + ctx.assertEquals(expected.toString(), text); + return (Void) null; + })); + })); + } + + @Test + public void testCopyOutAggregationThresholdProducesFewerChunks(TestContext ctx) { + run(ctx, conn -> conn.copyOut("COPY (SELECT generate_series(1, 500)) TO STDOUT", + new PgCopyOutOptions().setAggregationThreshold(4096)) + .compose(out -> { + AtomicInteger chunks = new AtomicInteger(); + Promise promise = Promise.promise(); + StringBuilder sb = new StringBuilder(); + out.handler(b -> { + chunks.incrementAndGet(); + sb.append(b.toString(StandardCharsets.UTF_8)); + }); + out.endHandler(v -> promise.complete(sb.toString())); + out.exceptionHandler(promise::tryFail); + return promise.future().map(text -> { + StringBuilder expected = new StringBuilder(); + for (int i = 1; i <= 500; i++) { + expected.append(i).append('\n'); + } + ctx.assertEquals(expected.toString(), text); + // 500 CopyData messages aggregated into a handful of buffers + ctx.assertTrue(chunks.get() < 100, "expected aggregation, got " + chunks.get() + " chunks"); + return null; + }); + })); + } + + @Test + public void testCopyOutExceptionHandlerSeesServerError(TestContext ctx) { + run(ctx, conn -> conn.copyOut("COPY (SELECT 1 / (5 - i) FROM generate_series(1, 10) AS i) TO STDOUT") + .compose(out -> { + Promise promise = Promise.promise(); + out.handler(b -> { + }); + out.exceptionHandler(promise::tryComplete); + out.endHandler(v -> promise.tryFail("expected a division by zero")); + return promise.future(); + }, err -> Future.succeededFuture(err)) + .map(err -> { + ctx.assertTrue(err instanceof PgException, "was " + err); + return null; + }) + // the connection stays usable after a failed COPY OUT + .compose(v -> conn.query("SELECT 1").execute()) + .mapEmpty()); + } + + @Test + public void testCopyInServerErrorIsReportedAndConnectionSurvives(TestContext ctx) { + run(ctx, conn -> table(conn) + .compose(v -> conn.copyIn(copyIn())) + // val is NOT NULL, so an empty value violates the constraint + .compose(in -> in.write(buf("1,alpha\nnot-an-int,beta\n")) + .compose(x -> in.end(), err -> Future.succeededFuture()) + .transform(ar -> in.completion()) + .otherwise(err -> { + ctx.assertTrue(err instanceof PgException, "was " + err); + return -1; + })) + .compose(v -> count(conn)) + .map(count -> { + ctx.assertEquals(0, count); + return null; + })); + } + + // ---------------------------------------------------------------- transactions + + @Test + public void testCopyInsideCommittedTransaction(TestContext ctx) { + run(ctx, conn -> table(conn) + .compose(v -> conn.begin()) + .compose(tx -> conn.copyIn(copyIn()) + .compose(in -> in.write(buf("1,alpha\n")).compose(x -> in.end())) + .compose(x -> tx.commit())) + .compose(v -> count(conn)) + .map(count -> { + ctx.assertEquals(1, count); + return null; + })); + } + + @Test + public void testCopyInsideRolledBackTransaction(TestContext ctx) { + run(ctx, conn -> table(conn) + .compose(v -> conn.begin()) + .compose(tx -> conn.copyIn(copyIn()) + .compose(in -> in.write(buf("1,alpha\n")).compose(x -> in.end())) + .compose(x -> tx.rollback())) + .compose(v -> count(conn)) + .map(count -> { + ctx.assertEquals(0, count); + return null; + })); + } + + // ---------------------------------------------------------------- pooling + + @Test + public void testPooledConnectionsRunCopiesConcurrently(TestContext ctx) { + Pool pool = PgBuilder.pool(b -> b.connectingTo(options).with(new PoolOptions().setMaxSize(4)).using(vertx)); + + List> results = new ArrayList<>(); + for (int i = 1; i <= 8; i++) { + int n = i; + results.add(pool.withConnection(conn -> ((PgConnection) conn) + .copyOut("COPY (SELECT generate_series(1, " + n + ")) TO STDOUT") + .compose(this::drain))); + } + + Future.all(new ArrayList<>(results)) + .map(composite -> { + for (int i = 1; i <= 8; i++) { + StringBuilder expected = new StringBuilder(); + for (int j = 1; j <= i; j++) { + expected.append(j).append('\n'); + } + ctx.assertEquals(expected.toString(), composite.resultAt(i - 1)); + } + return null; + }) + .eventually(pool::close) + .onComplete(ctx.asyncAssertSuccess()); + } + + @Test + public void testCopyOutAbandonedOnRecycleDoesNotWedgeThePool(TestContext ctx) { + Pool pool = PgBuilder.pool(b -> b.connectingTo(options).with(new PoolOptions().setMaxSize(1)).using(vertx)); + + pool.getConnection() + .compose(conn -> { + PgConnection pg = (PgConnection) conn; + // take the stream, consume nothing, hand the connection straight back + return pg.copyOut("COPY (SELECT generate_series(1, 2000)) TO STDOUT") + .compose(out -> { + out.pause(); + return pg.close(); + }); + }) + // the single pooled connection must be usable again + .compose(v -> pool.withConnection(conn -> conn.query("SELECT 42 AS n").execute())) + .map(rs -> { + ctx.assertEquals(42, rs.iterator().next().getInteger("n")); + return null; + }) + .eventually(pool::close) + .onComplete(ctx.asyncAssertSuccess()); + } + + @Test + public void testDefaultThresholdEmitsOneBufferPerCopyDataMessage(TestContext ctx) { + int rows = 500; + run(ctx, conn -> conn.copyOut( + "COPY (SELECT i, repeat('x', 40) FROM generate_series(1, " + rows + ") AS i) TO STDOUT WITH (FORMAT csv)", + new PgCopyOutOptions().setAggregationThreshold(1)) + .compose(out -> { + AtomicInteger chunks = new AtomicInteger(); + Promise promise = Promise.promise(); + out.handler(b -> chunks.incrementAndGet()); + out.endHandler(v -> promise.complete()); + out.exceptionHandler(promise::tryFail); + return promise.future().map(v -> chunks.get()); + }) + .map(chunks -> { + // PostgreSQL frames one CopyData message per row, and the default threshold of 1 hands + // each of them to the application as its own buffer + ctx.assertEquals(rows, chunks); + return null; + })); + } + + @Test + public void testConnectionLossDuringCopyInFailsCompletion(TestContext ctx) { + // The guarantee that matters: whatever happens to an individual write, a COPY that did not + // reach the server must be reported as failed by completion(). + Async async = ctx.async(); + PgConnection.connect(vertx, options).compose(victim -> + victim.query("SELECT pg_backend_pid() AS pid").execute().compose(rs -> { + int pid = rs.iterator().next().getInteger("pid"); + return victim.query("CREATE TEMP TABLE copy_loss (id INT, val TEXT)").execute() + .compose(v -> victim.copyIn("COPY copy_loss FROM STDIN WITH (FORMAT csv)", + new PgCopyInOptions())) + .compose(in -> in.write(buf("1,alpha\n")) + // kill the backend from another connection, then keep pushing data at a dead socket + .compose(v -> PgConnection.connect(vertx, options) + .compose(killer -> killer.query("SELECT pg_terminate_backend(" + pid + ")").execute() + .eventually(killer::close))) + .compose(v -> { + Future writes = Future.succeededFuture(); + for (int i = 0; i < 200; i++) { + int n = i; + writes = writes.transform(ar -> in.write(buf(n + ",value-" + n + "\n"))); + } + return writes.transform(ar -> in.end()).transform(ar -> in.completion()); + }) + .transform(ar -> { + ctx.assertTrue(ar.failed(), "completion() must fail when the connection dies"); + ctx.assertNotNull(ar.cause()); + return Future.succeededFuture(); + })); + }).eventually(victim::close) + ).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + @Test + public void testStatementThatNeverEntersCopyModeFailsRatherThanHangs(TestContext ctx) { + // The direction keyword sits inside a literal, so the client lets it through, but the server + // runs a COPY TO PROGRAM and never enters copy mode. There is no stream to hand out, and the + // caller must be told instead of waiting forever. + run(ctx, conn -> table(conn) + .compose(v -> assertFailsWithType(ctx, +conn.copyOut("COPY copy_contract TO PROGRAM 'cat > /dev/null # TO STDOUT'"), +IllegalStateException.class)) + // and the connection is still usable afterwards + .compose(v -> conn.query("SELECT 1").execute()) + .mapEmpty()); + } + + private static final String COPY_IN_SQL = "COPY copy_contract (id, val) FROM STDIN WITH (FORMAT csv)"; + + private static String copyIn() { + return COPY_IN_SQL; + } + + private void run(TestContext ctx, Function> body) { + PgConnection.connect(vertx, options) + .compose(conn -> { + Future op; + try { + op = body.apply(conn); + } catch (Throwable t) { + op = Future.failedFuture(t); + } + return op.eventually(conn::close); + }) + .onComplete(ctx.asyncAssertSuccess()); + } + + private static Future table(PgConnection conn) { + return exec(conn, "CREATE TEMP TABLE copy_contract (id INT PRIMARY KEY, val TEXT NOT NULL)"); + } + + private static Future copyRow(PgConnection conn, String csv) { + return conn.copyIn(COPY_IN_SQL) + .compose(in -> in.write(buf(csv)).compose(v -> in.end())) + .mapEmpty(); + } + + private static Future count(PgConnection conn) { + return conn.query("SELECT count(*) AS c FROM copy_contract").execute() + .map(rs -> rs.iterator().next().getInteger("c")); + } + + private static Future> rows(PgConnection conn) { + return conn.query("SELECT id, val FROM copy_contract ORDER BY id").execute() + .map(rs -> { + List out = new ArrayList<>(); + for (Row row : rs) { + out.add(row.getInteger("id") + ":" + row.getString("val")); + } + return out; + }); + } + + private Future drain(PgCopyOut out) { + return collectBytes(out).map(b -> b.toString(StandardCharsets.UTF_8)); + } + + private Future drain(PgCopyOut out, AtomicBoolean endHandlerFired) { + Promise promise = Promise.promise(); + Buffer acc = Buffer.buffer(); + out.handler(acc::appendBuffer); + out.endHandler(v -> { + endHandlerFired.set(true); + promise.tryComplete(acc); + }); + out.exceptionHandler(promise::tryFail); + return promise.future().map(b -> b.toString(StandardCharsets.UTF_8)); + } + + private Future collectBytes(PgCopyOut out) { + Promise promise = Promise.promise(); + Buffer acc = Buffer.buffer(); + out.handler(acc::appendBuffer); + out.endHandler(v -> promise.tryComplete(acc)); + out.exceptionHandler(promise::tryFail); + return promise.future(); + } + + private static Future assertFailsWithType(TestContext ctx, Future future, Class type) { + return future.compose( + v -> Future.failedFuture("expected a " + type.getSimpleName()), + err -> { + ctx.assertTrue(type.isInstance(err), "expected " + type.getSimpleName() + " but got " + err); + return Future.succeededFuture(); + }); + } + + private void assertCopyInCsv(TestContext ctx, PgCopyInOptions options, List expected) { + Async async = ctx.async(); + + withConnection(conn -> createTextTable(conn, "copy_test") + .compose(v -> conn.copyIn("COPY copy_test (id, val) FROM STDIN WITH (FORMAT csv)", options)) + .compose(in -> { + Future w1 = in.write(buf("1,alpha\n")); + Future w2 = in.write(buf("2,beta\n")); + Future w3 = in.write(buf("3,gamma\n")); + + Future completion = in.end(); + return Future.all(w1, w2, w3) + .compose(v -> completion) + .compose(v -> in.completion()) + .map(rowCount -> { + ctx.assertEquals(expected.size(), rowCount); + return null; + }); + }) + .compose(v -> fetchTextRows(conn, "copy_test")) + .map(rows -> { + ctx.assertEquals(expected, rows); + return null; + }) + ).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + private void assertCopyInSequentialComposedWrites(TestContext ctx) { + Async async = ctx.async(); + + withConnection(conn -> createTextTable(conn, "copy_test") + .compose(v -> conn.copyIn( + "COPY copy_test (id, val) FROM STDIN WITH (FORMAT csv)", + new PgCopyInOptions() + .setChunkSize(1024 * 1024))) + .compose(in -> in.write(buf("1,alpha\n")) + .compose(v -> in.write(buf("2,beta\n"))) + .compose(v -> in.write(buf("3,gamma\n"))) + .compose(v -> in.end())) + .compose(v -> fetchTextRows(conn, "copy_test")) + .map(rows -> { + ctx.assertEquals(Arrays.asList("1:alpha", "2:beta", "3:gamma"), rows); + return null; + }) + ).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + private void assertCopyInCsvSingleLargeBuffer(TestContext ctx, PgCopyInOptions options) { + Async async = ctx.async(); + + StringBuilder payload = new StringBuilder(); + List expected = new ArrayList<>(); + for (int i = 1; i <= 128; i++) { + payload.append(i).append(",value-").append(i).append('\n'); + expected.add(i + ":value-" + i); + } + + withConnection(conn -> createTextTable(conn, "copy_test") + .compose(v -> conn.copyIn("COPY copy_test (id, val) FROM STDIN WITH (FORMAT csv)", options)) + .compose(in -> { + Future write = in.write(buf(payload.toString())); + Future completion = in.end(); + return write.compose(v -> completion) + .compose(v -> in.completion()) + .map(rowCount -> { + ctx.assertEquals(expected.size(), rowCount); + return null; + }); + }) + .compose(v -> fetchTextRows(conn, "copy_test")) + .map(rows -> { + ctx.assertEquals(expected, rows); + return null; + }) + ).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + private void assertUnexpectedCopyQueryClosesConnection(TestContext ctx, String sql) { + Async async = ctx.async(); + + PgConnection.connect(vertx, options).compose(conn -> { + Promise closed = Promise.promise(); + AtomicReference exception = new AtomicReference<>(); + conn.exceptionHandler(exception::set); + conn.closeHandler(v -> closed.tryComplete()); + + return conn.query(sql).execute() + .compose(v -> Future.failedFuture("plain query unexpectedly handled COPY protocol"), err -> Future.succeededFuture()) + .compose(v -> closed.future()) + .map(v -> { + ctx.assertNotNull(exception.get()); + return null; + }); + }).onComplete(ctx.asyncAssertSuccess(v -> async.complete())); + } + + private Future unexpectedCopyQuery(PgConnection conn, String sql) { + Promise closed = Promise.promise(); + AtomicReference exception = new AtomicReference<>(); + conn.exceptionHandler(exception::set); + conn.closeHandler(v -> closed.tryComplete()); + + return conn.query(sql).execute() + .compose(v -> Future.failedFuture("plain query unexpectedly handled COPY protocol"), err -> Future.succeededFuture()) + .compose(v -> closed.future()) + .map(v -> { + if (exception.get() == null) { + throw new AssertionError("Expected protocol exception"); + } + return null; + }); + } + + private Future withConnection(Function> fn) { + Promise promise = Promise.promise(); + PgConnection.connect(vertx, options).onComplete(ar -> { + if (ar.failed()) { + promise.fail(ar.cause()); + return; + } + PgConnection conn = ar.result(); + Future op; + try { + op = fn.apply(conn); + } catch (Throwable t) { + op = Future.failedFuture(t); + } + op.eventually(conn::close).onComplete(promise); + }); + return promise.future(); + } + + private Future withTwoConnections(BiFunction> fn) { + Promise promise = Promise.promise(); + + PgConnection.connect(vertx, options).onComplete(srcAr -> { + if (srcAr.failed()) { + promise.fail(srcAr.cause()); + return; + } + + PgConnection src = srcAr.result(); + PgConnection.connect(vertx, options).onComplete(dstAr -> { + if (dstAr.failed()) { + src.close().onComplete(v -> promise.fail(dstAr.cause())); + return; + } + + PgConnection dst = dstAr.result(); + Future op; + try { + op = fn.apply(src, dst); + } catch (Throwable t) { + op = Future.failedFuture(t); + } + op.eventually(() -> src.close().eventually(dst::close)).onComplete(promise); + }); + }); + + return promise.future(); + } + + private static Future exec(PgConnection conn, String sql) { + return conn.query(sql).execute().mapEmpty(); + } + + private static Buffer buf(String value) { + return Buffer.buffer(value, StandardCharsets.UTF_8.name()); + } + + private static Future expectFailure(Future future, String message) { + return future.compose(v -> Future.failedFuture(message), err -> Future.succeededFuture()); + } + + private static Future assertFailsWith(TestContext ctx, Class expectedType, Future future) { + return future.compose( + v -> Future.failedFuture("Expected " + expectedType.getName()), + err -> { + ctx.assertTrue(expectedType.isInstance(err), "Expected " + expectedType.getName() + " but got " + err); + return Future.succeededFuture(); + }); + } + + private static void assertThrows(TestContext ctx, Class expectedType, Runnable action) { + try { + action.run(); + ctx.fail("Expected " + expectedType.getName()); + } catch (Throwable t) { + ctx.assertTrue(expectedType.isInstance(t)); + } + } + + private static Future createTextTable(PgConnection conn, String table) { + return exec(conn, "CREATE TEMP TABLE " + table + " (id INT PRIMARY KEY, val TEXT NOT NULL)"); + } + + private static Future insertTextRows(PgConnection conn, String table) { + return exec(conn, "INSERT INTO " + table + " (id, val) VALUES (1, 'alpha'), (2, 'beta'), (3, 'gamma')"); + } + + private static Future fetchCount(PgConnection conn, String table) { + return conn.query("SELECT count(*) AS c FROM " + table) + .execute() + .map(rows -> rows.iterator().next().getInteger("c")); + } + + private static ChannelConfig channelConfig(PgConnection conn) { + PgConnectionImpl connection = (PgConnectionImpl) conn; + PgSocketConnection socket = (PgSocketConnection) connection.unwrap().unwrap(); + return socket.socket().channelHandlerContext().channel().config(); + } + + private static void assertWatermark(TestContext ctx, WriteBufferWaterMark watermark, int low, int high) { + ctx.assertEquals(low, watermark.low()); + ctx.assertEquals(high, watermark.high()); + } + + private static Future> fetchTextRows(PgConnection conn, String table) { + return conn.query("SELECT id, val FROM " + table + " ORDER BY id") + .execute() + .map(rows -> { + List got = new ArrayList<>(); + for (Row row : rows) { + got.add(row.getInteger("id") + ":" + row.getString("val")); + } + return got; + }); + } + + private static Future collect(PgCopyOut out) { + Promise promise = Promise.promise(); + StringBuilder sb = new StringBuilder(); + + out.pause(); + out.handler(buffer -> sb.append(buffer.toString(StandardCharsets.UTF_8))); + out.exceptionHandler(promise::tryFail); + out.endHandler(v -> promise.tryComplete(sb.toString())); + out.completion().onFailure(promise::tryFail); + out.resume(); + + return promise.future(); + } + + private static Future collectFlowing(PgCopyOut out) { + Promise promise = Promise.promise(); + StringBuilder sb = new StringBuilder(); + + out.handler(buffer -> sb.append(buffer.toString(StandardCharsets.UTF_8))); + out.exceptionHandler(promise::tryFail); + out.endHandler(v -> promise.tryComplete(sb.toString())); + out.completion().onFailure(promise::tryFail); + + return promise.future(); + } + + private static String normalize(String value) { + return value.replace("\r\n", "\n"); + } + + +} diff --git a/vertx-pg-client/src/test/java/io/vertx/tests/pgclient/PgCopyTracingTest.java b/vertx-pg-client/src/test/java/io/vertx/tests/pgclient/PgCopyTracingTest.java new file mode 100644 index 000000000..e66f5a621 --- /dev/null +++ b/vertx-pg-client/src/test/java/io/vertx/tests/pgclient/PgCopyTracingTest.java @@ -0,0 +1,161 @@ +/* + * Copyright (c) 2011-2026 Contributors to the Eclipse Foundation + * + * This program and the accompanying materials are made available under the + * terms of the Eclipse Public License 2.0 which is available at + * http://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0 + * which is available at https://www.apache.org/licenses/LICENSE-2.0. + * + * SPDX-License-Identifier: EPL-2.0 OR Apache-2.0 + */ + +package io.vertx.tests.pgclient; + +import io.vertx.core.Context; +import io.vertx.core.Future; +import io.vertx.core.Promise; +import io.vertx.core.Vertx; +import io.vertx.core.VertxOptions; +import io.vertx.core.buffer.Buffer; +import io.vertx.core.spi.tracing.SpanKind; +import io.vertx.core.spi.tracing.TagExtractor; +import io.vertx.core.spi.tracing.VertxTracer; +import io.vertx.core.tracing.TracingOptions; +import io.vertx.core.tracing.TracingPolicy; +import io.vertx.ext.unit.TestContext; +import io.vertx.ext.unit.junit.VertxUnitRunner; +import io.vertx.pgclient.PgConnection; +import io.vertx.pgclient.PgCopyOut; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; + +import java.nio.charset.StandardCharsets; +import java.util.List; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.function.BiConsumer; +import java.util.function.Function; + +/** + * COPY statements must show up in traces like any other statement the application runs. + */ +@RunWith(VertxUnitRunner.class) +public class PgCopyTracingTest extends PgTestBase { + + private Vertx vertx; + private final List requests = new CopyOnWriteArrayList<>(); + private final List responses = new CopyOnWriteArrayList<>(); + + @Before + public void setup() throws Exception { + super.setup(); + requests.clear(); + responses.clear(); + VertxTracer tracer = new VertxTracer() { + @Override + public Object sendRequest(Context context, SpanKind kind, TracingPolicy policy, R request, + String operation, BiConsumer headers, TagExtractor extractor) { + requests.add(operation); + return request; + } + + @Override + public void receiveResponse(Context context, R response, Object payload, Throwable failure, + TagExtractor extractor) { + responses.add(failure == null ? "ok" : "failed"); + } + }; + vertx = Vertx.builder() + .with(new VertxOptions().setTracingOptions(new TracingOptions())) + .withTracer(options -> tracer) + .build(); + } + + @After + public void tearDown(TestContext ctx) { + vertx.close().onComplete(ctx.asyncAssertSuccess()); + } + + @Test + public void testCopyInIsTraced(TestContext ctx) { + run(ctx, conn -> conn.query("CREATE TEMP TABLE copy_trace (id INT, val TEXT)").execute() + .compose(v -> { + requests.clear(); + responses.clear(); + return conn.copyIn("COPY copy_trace FROM STDIN WITH (FORMAT csv)"); + }) + .compose(in -> in.write(Buffer.buffer("1,alpha\n", StandardCharsets.UTF_8.name())) + .compose(x -> in.end()) + .compose(x -> in.completion())) + .map(rows -> { + ctx.assertEquals(1, rows); + ctx.assertEquals(1, requests.size(), "COPY IN produced no trace span, spans=" + requests); + ctx.assertEquals("Query", requests.get(0)); + ctx.assertEquals(1, responses.size()); + ctx.assertEquals("ok", responses.get(0)); + return null; + })); + } + + @Test + public void testCopyOutIsTraced(TestContext ctx) { + run(ctx, conn -> { + requests.clear(); + responses.clear(); + return conn.copyOut("COPY (SELECT generate_series(1, 3)) TO STDOUT") + .compose(out -> drain(out).compose(text -> out.completion())) + .map(rows -> { + ctx.assertEquals(3, rows); + ctx.assertEquals(1, requests.size(), "COPY OUT produced no trace span, spans=" + requests); + ctx.assertEquals("Query", requests.get(0)); + ctx.assertEquals(1, responses.size()); + ctx.assertEquals("ok", responses.get(0)); + return null; + }); + }); + } + + @Test + public void testFailedCopyIsReportedAsAFailedSpan(TestContext ctx) { + run(ctx, conn -> { + requests.clear(); + responses.clear(); + return conn.copyOut("COPY (SELECT 1 / 0) TO STDOUT") + .compose(out -> drain(out).mapEmpty(), err -> Future.succeededFuture()) + // the span closes at ReadyForQuery, which is also what lets the next statement run + .compose(v -> conn.query("SELECT 1").execute()) + .map(v -> { + // the trailing SELECT is traced too, so assert on the first span, which is the COPY + ctx.assertFalse(requests.isEmpty(), "no span for the failed COPY"); + ctx.assertEquals("Query", requests.get(0)); + ctx.assertFalse(responses.isEmpty(), "the COPY span was never closed"); + ctx.assertEquals("failed", responses.get(0)); + return null; + }); + }); + } + + private void run(TestContext ctx, Function> body) { + PgConnection.connect(vertx, options) + .compose(conn -> { + Future op; + try { + op = body.apply(conn); + } catch (Throwable t) { + op = Future.failedFuture(t); + } + return op.eventually(conn::close); + }) + .onComplete(ctx.asyncAssertSuccess()); + } + + private Future drain(PgCopyOut out) { + Promise promise = Promise.promise(); + Buffer acc = Buffer.buffer(); + out.handler(acc::appendBuffer); + out.endHandler(v -> promise.tryComplete(acc)); + out.exceptionHandler(promise::tryFail); + return promise.future().map(b -> b.toString(StandardCharsets.UTF_8)); + } +} diff --git a/vertx-sql-client/src/main/java/io/vertx/sqlclient/impl/pool/SqlConnectionPool.java b/vertx-sql-client/src/main/java/io/vertx/sqlclient/impl/pool/SqlConnectionPool.java index b3c2d7b87..36f573ef4 100644 --- a/vertx-sql-client/src/main/java/io/vertx/sqlclient/impl/pool/SqlConnectionPool.java +++ b/vertx-sql-client/src/main/java/io/vertx/sqlclient/impl/pool/SqlConnectionPool.java @@ -30,7 +30,7 @@ import io.vertx.sqlclient.spi.connection.ConnectionContext; import io.vertx.sqlclient.spi.connection.ConnectionFactory; import io.vertx.sqlclient.spi.protocol.CommandBase; -import io.vertx.sqlclient.spi.protocol.QueryCommandBase; +import io.vertx.sqlclient.spi.protocol.SqlCommand; import java.util.List; import java.util.function.Function; @@ -428,8 +428,8 @@ public void schedule(CommandBase cmd, Completable handler) { QueryReporter queryReporter; VertxTracer tracer = vertx.tracer(); ClientMetrics metrics = conn.metrics(); - if (cmd instanceof QueryCommandBase && (tracer != null || metrics != null)) { - queryReporter = new QueryReporter(tracer, metrics, context, (QueryCommandBase) cmd, conn); + if (cmd instanceof SqlCommand && (tracer != null || metrics != null)) { + queryReporter = new QueryReporter(tracer, metrics, context, cmd, conn); queryReporter.before(); } else { queryReporter = null; diff --git a/vertx-sql-client/src/main/java/io/vertx/sqlclient/impl/tracing/QueryReporter.java b/vertx-sql-client/src/main/java/io/vertx/sqlclient/impl/tracing/QueryReporter.java index 47d699dc7..fbcae6175 100644 --- a/vertx-sql-client/src/main/java/io/vertx/sqlclient/impl/tracing/QueryReporter.java +++ b/vertx-sql-client/src/main/java/io/vertx/sqlclient/impl/tracing/QueryReporter.java @@ -10,8 +10,9 @@ import io.vertx.sqlclient.impl.QueryResultBuilder; import io.vertx.sqlclient.spi.connection.Connection; import io.vertx.sqlclient.spi.protocol.ExtendedQueryCommand; +import io.vertx.sqlclient.spi.protocol.CommandBase; import io.vertx.sqlclient.spi.protocol.QueryCommandBase; -import io.vertx.sqlclient.spi.protocol.SimpleQueryCommand; +import io.vertx.sqlclient.spi.protocol.SqlCommand; import java.util.Collections; import java.util.List; @@ -66,7 +67,7 @@ public String value(QueryRequest obj, int index) { } }; - private final QueryCommandBase cmd; + private final CommandBase cmd; private final VertxTracer tracer; private final ClientMetrics metrics; private final ContextInternal context; @@ -78,7 +79,7 @@ public String value(QueryRequest obj, int index) { private Object payload; private Object metric; - public QueryReporter(VertxTracer tracer, ClientMetrics metrics, ContextInternal context, QueryCommandBase queryCmd, Connection conn) { + public QueryReporter(VertxTracer tracer, ClientMetrics metrics, ContextInternal context, CommandBase queryCmd, Connection conn) { this.tracer = tracer; this.metrics = metrics; this.context = context; @@ -109,22 +110,27 @@ private void receiveResponse(ContextInternal context, Object payload, Object res tracer.receiveResponse(context, result, payload, failure, TagExtractor.empty()); } + private String sql() { + return ((SqlCommand) cmd).sql(); + } + public void before() { if (tracer != null) { - String sql = cmd.sql(); - if (cmd instanceof SimpleQueryCommand) { - payload = sendRequest(context, sql); - } else { + String sql = sql(); + if (cmd instanceof ExtendedQueryCommand) { ExtendedQueryCommand extendedQueryCmd = (ExtendedQueryCommand) cmd; if (extendedQueryCmd.params() != null) { - payload = sendRequest(context, sql, ((ExtendedQueryCommand) cmd).params()); + payload = sendRequest(context, sql, extendedQueryCmd.params()); } else { payload = sendRequest(context, sql, ((ExtendedQueryCommand) cmd).paramsList()); } + } else { + // simple queries and statements that carry no parameters, such as COPY + payload = sendRequest(context, sql); } } if (metrics != null) { - String sql = cmd.sql(); + String sql = sql(); metric = metrics.init(); metrics.requestBegin(metric, sql, sql); metrics.requestEnd(metric); @@ -133,8 +139,12 @@ public void before() { public void after(Object res, Throwable err) { if (tracer != null) { - QueryResultBuilder qbr = (QueryResultBuilder) cmd.resultHandler(); - receiveResponse(context, payload, err == null ? qbr.first : null, err); + Object result = null; + if (err == null && cmd instanceof QueryCommandBase) { + QueryResultBuilder qbr = (QueryResultBuilder) ((QueryCommandBase) cmd).resultHandler(); + result = qbr.first; + } + receiveResponse(context, payload, result, err); } if (metrics != null) { if (err == null) { diff --git a/vertx-sql-client/src/main/java/io/vertx/sqlclient/internal/SqlConnectionBase.java b/vertx-sql-client/src/main/java/io/vertx/sqlclient/internal/SqlConnectionBase.java index b68418703..1a1b82ec3 100644 --- a/vertx-sql-client/src/main/java/io/vertx/sqlclient/internal/SqlConnectionBase.java +++ b/vertx-sql-client/src/main/java/io/vertx/sqlclient/internal/SqlConnectionBase.java @@ -31,7 +31,7 @@ import io.vertx.sqlclient.spi.protocol.CommandBase; import io.vertx.sqlclient.spi.connection.Connection; import io.vertx.sqlclient.spi.protocol.PrepareStatementCommand; -import io.vertx.sqlclient.spi.protocol.QueryCommandBase; +import io.vertx.sqlclient.spi.protocol.SqlCommand; import io.vertx.sqlclient.impl.pool.SqlConnectionPool; import io.vertx.sqlclient.impl.tracing.QueryReporter; import io.vertx.sqlclient.spi.connection.ConnectionFactory; @@ -162,8 +162,8 @@ public void schedule(CommandBase cmd, Completable handler) { QueryReporter queryReporter; VertxTracer tracer = context.owner().tracer(); ClientMetrics metrics = conn.metrics(); - if (!(conn instanceof SqlConnectionPool.PooledConnection) && cmd instanceof QueryCommandBase && (tracer != null || metrics != null)) { - queryReporter = new QueryReporter(tracer, metrics, context, (QueryCommandBase) cmd, conn); + if (!(conn instanceof SqlConnectionPool.PooledConnection) && cmd instanceof SqlCommand && (tracer != null || metrics != null)) { + queryReporter = new QueryReporter(tracer, metrics, context, cmd, conn); queryReporter.before(); conn .schedule(cmd, (res, err) -> { diff --git a/vertx-sql-client/src/main/java/io/vertx/sqlclient/spi/protocol/QueryCommandBase.java b/vertx-sql-client/src/main/java/io/vertx/sqlclient/spi/protocol/QueryCommandBase.java index 9a02e4758..f4b3df1cb 100644 --- a/vertx-sql-client/src/main/java/io/vertx/sqlclient/spi/protocol/QueryCommandBase.java +++ b/vertx-sql-client/src/main/java/io/vertx/sqlclient/spi/protocol/QueryCommandBase.java @@ -26,7 +26,7 @@ * @author Julien Viet */ -public abstract class QueryCommandBase extends CommandBase { +public abstract class QueryCommandBase extends CommandBase implements SqlCommand { private final QueryResultHandler resultHandler; private final Collector collector; diff --git a/vertx-sql-client/src/main/java/io/vertx/sqlclient/spi/protocol/SqlCommand.java b/vertx-sql-client/src/main/java/io/vertx/sqlclient/spi/protocol/SqlCommand.java new file mode 100644 index 000000000..3104f5b79 --- /dev/null +++ b/vertx-sql-client/src/main/java/io/vertx/sqlclient/spi/protocol/SqlCommand.java @@ -0,0 +1,26 @@ +/* + * Copyright (c) 2011-2026 Contributors to the Eclipse Foundation + * + * This program and the accompanying materials are made available under the + * terms of the Eclipse Public License 2.0 which is available at + * http://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0 + * which is available at https://www.apache.org/licenses/LICENSE-2.0. + * + * SPDX-License-Identifier: EPL-2.0 OR Apache-2.0 + */ + +package io.vertx.sqlclient.spi.protocol; + +/** + * A command that runs a statement written by the application. + *

+ * Implemented by queries and by any other command carrying user SQL, such as COPY. Tracing and + * metrics report these and ignore the rest. + */ +public interface SqlCommand { + + /** + * @return the statement sent to the database + */ + String sql(); +}