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(); +}