diff --git a/src/main/java/io/r2dbc/postgresql/PostgresqlConnection.java b/src/main/java/io/r2dbc/postgresql/PostgresqlConnection.java index 1686b3d2..239f57a1 100644 --- a/src/main/java/io/r2dbc/postgresql/PostgresqlConnection.java +++ b/src/main/java/io/r2dbc/postgresql/PostgresqlConnection.java @@ -16,6 +16,8 @@ package io.r2dbc.postgresql; +import io.netty.util.ReferenceCountUtil; +import io.netty.util.ReferenceCounted; import io.r2dbc.postgresql.api.CopyInBuilder; import io.r2dbc.postgresql.api.ErrorDetails; import io.r2dbc.postgresql.api.Notification; @@ -58,6 +60,7 @@ import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; +import static io.r2dbc.postgresql.PostgresqlStatement.WINDOW_UNTIL; import static io.r2dbc.postgresql.client.TransactionStatus.IDLE; import static io.r2dbc.postgresql.client.TransactionStatus.OPEN; @@ -233,6 +236,17 @@ public CopyInBuilder copyIn(String sql) { return new PostgresqlCopyIn.Builder(this.resources, sql); } + @Override + public Flux copyOut(String sql) { + ExceptionFactory factory = ExceptionFactory.withSql(sql); + + return SimpleQueryMessageFlow.exchange(this.resources.getClient(), sql) + .windowUntil(WINDOW_UNTIL) + .doOnDiscard(ReferenceCounted .class, ReferenceCountUtil::release) // ensure release of rows within WindowPredicate + .map(messages -> PostgresqlCopyOutResult.toCopyOutResult(messages, factory)); + + } + @Override public PostgresqlBatch createBatch() { return new PostgresqlBatch(this.resources); diff --git a/src/main/java/io/r2dbc/postgresql/PostgresqlCopyOutResult.java b/src/main/java/io/r2dbc/postgresql/PostgresqlCopyOutResult.java new file mode 100644 index 00000000..b4c2d6d9 --- /dev/null +++ b/src/main/java/io/r2dbc/postgresql/PostgresqlCopyOutResult.java @@ -0,0 +1,152 @@ +/* + * Copyright 2019 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.r2dbc.postgresql; + +import io.netty.buffer.ByteBuf; +import io.netty.util.AbstractReferenceCounted; +import io.netty.util.ReferenceCountUtil; +import io.netty.util.ReferenceCounted; +import io.r2dbc.postgresql.message.Format; +import io.r2dbc.postgresql.message.backend.*; +import io.r2dbc.postgresql.util.Assert; +import reactor.core.publisher.Flux; +import reactor.core.publisher.SynchronousSink; + +import java.util.Collection; +import java.util.Objects; +import java.util.function.BiFunction; + +public class PostgresqlCopyOutResult extends AbstractReferenceCounted implements io.r2dbc.postgresql.api.PostgresqlCopyOutResult { + + private final Flux messages; + + private final ExceptionFactory factory; + + private volatile CopyOutMetadata metadata; + + PostgresqlCopyOutResult(Flux messages, ExceptionFactory factory) { + this.messages = Assert.requireNonNull(messages, "messages must not be null"); + this.factory = Assert.requireNonNull(factory, "factory must not be null"); + } + + + @Override + @SuppressWarnings({"rawtypes", "unchecked"}) + public Flux map(BiFunction f) { + Assert.requireNonNull(f, "f must not be null"); + + return (Flux) this.messages + .handle((message, sink) -> { + + try { + if (message instanceof ErrorResponse) { + this.factory.handleErrorResponse(message, (SynchronousSink) sink); + return; + } + + if (message instanceof CopyOutResponse) { + this.metadata = PostgresCopyOutMetadata.toMetadata((CopyOutResponse) message); + return; + } + + if (message instanceof CopyData) { + CopyData data = (CopyData) message; + sink.next(f.apply(data.getData(), this.metadata)); + return; + } + + if (message instanceof CopyDone) { + sink.complete(); + return; + } + + if (message instanceof RowDescription) { + sink.error(new IllegalStateException("copyOut may only be used with statements that execute a COPY ... TO STDOUT query")); + } + + } finally { + ReferenceCountUtil.release(message); + } + }); + } + + @Override + protected void deallocate() { + // drain messages for cleanup + messages.map(ReferenceCountUtil::release).subscribe(); + } + + @Override + public ReferenceCounted touch(Object hint) { + return this; + } + + @Override + public String toString() { + return "PostgresqlCopyOutResult{" + + "messages=" + messages + + ", factory=" + factory + + '}'; + } + + static PostgresqlCopyOutResult toCopyOutResult(Flux messages, ExceptionFactory factory) { + return new PostgresqlCopyOutResult(messages, factory); + } + + private static class PostgresCopyOutMetadata implements CopyOutMetadata { + private final Format overallFormat; + private final Collection columnFormats; + + private PostgresCopyOutMetadata(Format overallFormat, Collection columnFormats) { + this.overallFormat = overallFormat; + this.columnFormats = columnFormats; + } + + @Override + public Format getOverallFormat() { + return overallFormat; + } + + public Collection getColumnFormats() { + return columnFormats; + } + + @Override + public String toString() { + return "PostgresCopyOutMetadata{" + + "overallFormat=" + overallFormat + + ", columnFormats=" + columnFormats + + '}'; + } + + @Override + public boolean equals(Object o) { + if (!(o instanceof CopyOutMetadata)) return false; + CopyOutMetadata that = (CopyOutMetadata) o; + return overallFormat == that.getOverallFormat() && Objects.equals(columnFormats, that.getColumnFormats()); + } + + @Override + public int hashCode() { + return Objects.hash(overallFormat, columnFormats); + } + + static CopyOutMetadata toMetadata(CopyOutResponse response) { + return new PostgresCopyOutMetadata(response.getOverallFormat(), response.getColumnFormats()); + } + } +} diff --git a/src/main/java/io/r2dbc/postgresql/PostgresqlResult.java b/src/main/java/io/r2dbc/postgresql/PostgresqlResult.java index ea7ce1df..a25076f2 100644 --- a/src/main/java/io/r2dbc/postgresql/PostgresqlResult.java +++ b/src/main/java/io/r2dbc/postgresql/PostgresqlResult.java @@ -19,11 +19,7 @@ import io.netty.util.AbstractReferenceCounted; import io.netty.util.ReferenceCountUtil; import io.netty.util.ReferenceCounted; -import io.r2dbc.postgresql.message.backend.BackendMessage; -import io.r2dbc.postgresql.message.backend.CommandComplete; -import io.r2dbc.postgresql.message.backend.DataRow; -import io.r2dbc.postgresql.message.backend.ErrorResponse; -import io.r2dbc.postgresql.message.backend.RowDescription; +import io.r2dbc.postgresql.message.backend.*; import io.r2dbc.postgresql.util.Assert; import io.r2dbc.spi.Result; import io.r2dbc.spi.Row; @@ -127,6 +123,16 @@ public Flux map(BiFunction f) { if (message instanceof DataRow) { PostgresqlRow row = PostgresqlRow.toRow(this.resources, (DataRow) message, this.metadata, this.rowDescription); sink.next(f.apply(row, this.metadata)); + return; + } + + if (message instanceof CopyOutResponse) { + sink.error(new IllegalStateException("COPY ... TO STDOUT statements may only be used with PostgresqlConnection.copyOut()")); + return; + } + + if (message instanceof CopyInResponse) { + sink.error(new IllegalStateException("COPY ... FROM STDIN statements may only be used with PostgresqlConnection.copyIn()")); } } finally { diff --git a/src/main/java/io/r2dbc/postgresql/PostgresqlStatement.java b/src/main/java/io/r2dbc/postgresql/PostgresqlStatement.java index b714e711..cfbf8b13 100644 --- a/src/main/java/io/r2dbc/postgresql/PostgresqlStatement.java +++ b/src/main/java/io/r2dbc/postgresql/PostgresqlStatement.java @@ -57,7 +57,7 @@ */ final class PostgresqlStatement implements io.r2dbc.postgresql.api.PostgresqlStatement { - private static final Predicate WINDOW_UNTIL = or(CommandComplete.class::isInstance, EmptyQueryResponse.class::isInstance, ErrorResponse.class::isInstance); + static final Predicate WINDOW_UNTIL = or(CommandComplete.class::isInstance, EmptyQueryResponse.class::isInstance, ErrorResponse.class::isInstance); private final ArrayDeque bindings; diff --git a/src/main/java/io/r2dbc/postgresql/api/PostgresqlConnection.java b/src/main/java/io/r2dbc/postgresql/api/PostgresqlConnection.java index a1309e46..df9c20e6 100644 --- a/src/main/java/io/r2dbc/postgresql/api/PostgresqlConnection.java +++ b/src/main/java/io/r2dbc/postgresql/api/PostgresqlConnection.java @@ -18,11 +18,7 @@ import io.netty.buffer.ByteBuf; import io.r2dbc.postgresql.message.frontend.CancelRequest; -import io.r2dbc.spi.Connection; -import io.r2dbc.spi.IsolationLevel; -import io.r2dbc.spi.R2dbcNonTransientResourceException; -import io.r2dbc.spi.TransactionDefinition; -import io.r2dbc.spi.ValidationDepth; +import io.r2dbc.spi.*; import org.reactivestreams.Publisher; import org.reactivestreams.Subscriber; import reactor.core.publisher.Flux; @@ -70,6 +66,15 @@ public interface PostgresqlConnection extends Connection { */ CopyInBuilder copyIn(String sql); + /** + * Executes one or more PostgreSQL COPY ... TO STDOUT statements and returns the {@link PostgresqlCopyOutResult}s. + * {@link PostgresqlCopyOutResult} objects must be fully consumed to ensure full execution of the {@link Statement}. + * + * @return the {@link PostgresqlCopyOutResult}s, returned by each statement + * @since 1.2 + */ + Flux copyOut(String sql); + /** * Use {@code COPY FROM STDIN} for very fast copying into a database table. * diff --git a/src/main/java/io/r2dbc/postgresql/api/PostgresqlCopyOutResult.java b/src/main/java/io/r2dbc/postgresql/api/PostgresqlCopyOutResult.java new file mode 100644 index 00000000..81538cab --- /dev/null +++ b/src/main/java/io/r2dbc/postgresql/api/PostgresqlCopyOutResult.java @@ -0,0 +1,70 @@ +/* + * Copyright 2019 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.r2dbc.postgresql.api; + +import io.netty.buffer.ByteBuf; +import io.netty.util.ReferenceCounted; +import io.r2dbc.postgresql.message.Format; +import reactor.core.publisher.Flux; + +import java.util.Collection; +import java.util.function.BiFunction; + +/** + * A result obtained by executing a {@code COPY ... TO STDOUT} statement. + */ +public interface PostgresqlCopyOutResult { + + /** + * Returns a mapping of the rows that are the results of a copy-out operation against a database. May be empty if + * the query did not return any rows. + *

According to the PostgreSQL documentation, each {@link ByteBuf} will represent a complete row either in + * binary or textual representation, however it is not clear whether that is just an incidental implementation + * quirk or a strict guarantee. In the other direction, there is no such requirement, so robust consumers should + * probably not make any assumptions here. + *

The {@link ByteBuf} instances are only valid to use during the {@code mappingFunction} callback, unless they + * are retained (via {@link ReferenceCounted#retain()} and later released. When retaining them, care must be taken + * to ensure that they do eventually get released, as the asynchronous nature of streams can often see messages + * dropped when the stream is canceled downstream. + * + * @param mappingFunction the {@link BiFunction} that maps a {@link ByteBuf} and {@link CopyOutMetadata} to a value + * @param the type of the mapped value + * @return a mapping of the rows that are the results of a copy-out operation against a database + * @throws IllegalArgumentException if {@code mappingFunction} is {@code null} + * @throws IllegalStateException if the result was consumed + */ + Flux map(BiFunction mappingFunction); + + /** + * Metadata for a copy-out operation. + */ + interface CopyOutMetadata { + /** + * Get the format of each of the columns. + * + * @return The format of each of the columns + */ + Collection getColumnFormats(); + + /** + * Get the overall format of a response. + * + * @return The format. + */ + Format getOverallFormat(); + } +} diff --git a/src/test/java/io/r2dbc/postgresql/PostgresqlCopyOutResultIntegrationTests.java b/src/test/java/io/r2dbc/postgresql/PostgresqlCopyOutResultIntegrationTests.java new file mode 100644 index 00000000..f28ecef3 --- /dev/null +++ b/src/test/java/io/r2dbc/postgresql/PostgresqlCopyOutResultIntegrationTests.java @@ -0,0 +1,118 @@ +/* + * Copyright 2022 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.r2dbc.postgresql; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; +import io.r2dbc.postgresql.ExceptionFactory.PostgresqlBadGrammarException; +import io.r2dbc.spi.R2dbcNonTransientResourceException; +import org.awaitility.Awaitility; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.springframework.jdbc.core.JdbcOperations; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; +import reactor.test.StepVerifier; + +import java.nio.charset.Charset; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.Collections; +import java.util.List; + +import static java.util.Arrays.asList; +import static java.util.Collections.emptyList; +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Integration tests for {@link PostgresqlCopyOutResult}. + */ +class PostgresqlCopyOutResultIntegrationTests extends AbstractIntegrationTests { + + @BeforeEach + void setUp() { + super.setUp(); + getJdbcOperations().execute("DROP TABLE IF EXISTS test"); + getJdbcOperations().execute("CREATE TABLE test (id INTEGER PRIMARY KEY, value VARCHAR(255))"); + + } + + @AfterEach + void tearDown() { + super.tearDown(); + getJdbcOperations().execute("DROP TABLE IF EXISTS test"); + } + + private JdbcOperations getJdbcOperations() { + return SERVER.getJdbcOperations(); + } + + @Override + protected void customize(PostgresqlConnectionConfiguration.Builder builder) { + builder.preparedStatementCacheQueries(2); + } + + @Test + void shouldCopyDataOutOfTable() { + getJdbcOperations().execute("INSERT INTO test (id, value) VALUES (1, 'a'), (2, 'b'), (3, 'c')"); + + String sql = "COPY (SELECT id, value FROM test ORDER BY id) TO STDOUT WITH CSV"; + + this.connection.copyOut(sql) + .flatMap(result -> result.map(this::convertRow)) + .collectSortedList() + .as(StepVerifier::create) + .expectNext(asList("1,a\n", "2,b\n", "3,c\n")) + .verifyComplete(); + } + + @Test + void shouldCopyDataOutOfEmptyTable() { + String sql = "COPY (SELECT id, value FROM test ORDER BY id) TO STDOUT WITH CSV"; + + this.connection.copyOut(sql) + .flatMap(result -> result.map(this::convertRow)) + .as(StepVerifier::create) + .verifyComplete(); + } + + @Test + void shouldFailOnInvalidStatement() { + String sql = "COPY invalid command TO STDOUT"; + + this.connection.copyOut(sql) + .flatMap(result -> result.map(this::convertRow)) + .as(StepVerifier::create) + .verifyError(); + } + + @Test + void shouldFailOnOrdinarySelectStatement() { + getJdbcOperations().execute("INSERT INTO test (id, value) VALUES (1, 'a'), (2, 'b'), (3, 'c')"); + String sql = "SELECT * FROM test"; + + this.connection.copyOut(sql) + .flatMap(result -> result.map(this::convertRow)) + .as(StepVerifier::create) + .verifyErrorMatches(e -> e.getMessage().contains("copyOut may only be used")); + } + + private String convertRow(ByteBuf bytes, io.r2dbc.postgresql.api.PostgresqlCopyOutResult.CopyOutMetadata meta) { + return bytes.toString(StandardCharsets.UTF_8); + } +} diff --git a/src/test/java/io/r2dbc/postgresql/PostgresqlCopyOutResultUnitTests.java b/src/test/java/io/r2dbc/postgresql/PostgresqlCopyOutResultUnitTests.java new file mode 100644 index 00000000..a2077b01 --- /dev/null +++ b/src/test/java/io/r2dbc/postgresql/PostgresqlCopyOutResultUnitTests.java @@ -0,0 +1,69 @@ +/* + * Copyright 2017 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.r2dbc.postgresql; + +import io.netty.buffer.Unpooled; +import io.r2dbc.postgresql.message.Format; +import io.r2dbc.postgresql.message.backend.*; +import org.junit.jupiter.api.Test; +import reactor.core.publisher.Flux; +import reactor.test.StepVerifier; + +import java.util.Collections; + +import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException; + +/** + * Unit tests for {@link PostgresqlCopyOutResult}. + */ +final class PostgresqlCopyOutResultUnitTests { + + @Test + void toCopyOutResultCommandComplete() { + PostgresqlCopyOutResult result = PostgresqlCopyOutResult.toCopyOutResult(Flux.just(new CommandComplete("test", null, 0L)), ExceptionFactory.INSTANCE); + + result.map((row, rowMetadata) -> row) + .as(StepVerifier::create) + .verifyComplete(); + } + + @Test + void toCopyOutResultNoMessages() { + assertThatIllegalArgumentException().isThrownBy(() -> PostgresqlCopyOutResult.toCopyOutResult(null, ExceptionFactory.INSTANCE)) + .withMessage("messages must not be null"); + } + + @Test + void toCopyOutResultResponseMetadataMap() { + PostgresqlCopyOutResult result = PostgresqlCopyOutResult.toCopyOutResult(Flux.just(new CopyOutResponse(Collections.emptyList(), Format.FORMAT_BINARY), new CopyData(Unpooled.EMPTY_BUFFER), CopyDone.INSTANCE, new CommandComplete("test", null, null)), ExceptionFactory.INSTANCE); + + result.map((row, rowMetadata) -> row) + .as(StepVerifier::create) + .expectNextCount(1) + .verifyComplete(); + } + + @Test + void toCopyOutResultContainsRowDescription() { + io.r2dbc.postgresql.api.PostgresqlCopyOutResult result = PostgresqlCopyOutResult.toCopyOutResult(Flux.just(new RowDescription(Collections.emptyList()), new CommandComplete("test", null, null)), ExceptionFactory.INSTANCE); + + result.map((row, rowMetadata) -> row) + .as(StepVerifier::create) + .verifyErrorMatches(e -> e.getMessage().contains("copyOut may only be used")); + } + +} diff --git a/src/test/java/io/r2dbc/postgresql/api/MockPostgresqlConnection.java b/src/test/java/io/r2dbc/postgresql/api/MockPostgresqlConnection.java index 1d7709ee..88ddfb65 100644 --- a/src/test/java/io/r2dbc/postgresql/api/MockPostgresqlConnection.java +++ b/src/test/java/io/r2dbc/postgresql/api/MockPostgresqlConnection.java @@ -57,6 +57,11 @@ public CopyInBuilder copyIn(String sql) { throw new UnsupportedOperationException(); } + @Override + public Flux copyOut(String sql) { + return Flux.empty(); + } + @Override public PostgresqlBatch createBatch() { return null;