Stream for subscribing to Postgres connection notices

Francesco Komauli · Francesco Komauli · commit b8502da423b9 · 2022-12-02T10:26:37.000+01:00
[resolves pgjdbc#570] Postgres may send notice messages through a connection, which may contain log information or metadata related to the submitted commands. Notices travel on the same backend messaging subsystem used by notifications. The implementation processes NoticeResponse in a way similar to NotificationResponse, and exposes a coherent API.
diff --git a/src/main/java/io/r2dbc/postgresql/PostgresqlConnection.java b/src/main/java/io/r2dbc/postgresql/PostgresqlConnection.java
@@ -18,6 +18,7 @@
 
 import io.r2dbc.postgresql.api.CopyInBuilder;
 import io.r2dbc.postgresql.api.ErrorDetails;
+import io.r2dbc.postgresql.api.Notice;
 import io.r2dbc.postgresql.api.Notification;
 import io.r2dbc.postgresql.api.PostgresTransactionDefinition;
 import io.r2dbc.postgresql.api.PostgresqlResult;
@@ -31,6 +32,8 @@
 import io.r2dbc.postgresql.message.backend.BackendMessage;
 import io.r2dbc.postgresql.message.backend.CommandComplete;
 import io.r2dbc.postgresql.message.backend.ErrorResponse;
+import io.r2dbc.postgresql.message.backend.Field;
+import io.r2dbc.postgresql.message.backend.NoticeResponse;
 import io.r2dbc.postgresql.message.backend.NotificationResponse;
 import io.r2dbc.postgresql.util.Assert;
 import io.r2dbc.postgresql.util.Operators;
@@ -53,6 +56,8 @@
 import reactor.util.annotation.Nullable;
 
 import java.time.Duration;
+import java.util.EnumMap;
+import java.util.Map;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
 
@@ -78,6 +83,8 @@ final class PostgresqlConnection implements io.r2dbc.postgresql.api.PostgresqlCo
 
     private final AtomicReference<NotificationAdapter> notificationAdapter = new AtomicReference<>();
 
+    private final AtomicReference<NoticeAdapter> noticeAdapter = new AtomicReference<>();
+
     private volatile IsolationLevel isolationLevel;
 
     private volatile IsolationLevel previousIsolationLevel;
@@ -181,6 +188,12 @@ public Mono<Void> close() {
             if (notificationAdapter != null && this.notificationAdapter.compareAndSet(notificationAdapter, null)) {
                 notificationAdapter.dispose();
             }
+
+            NoticeAdapter noticeAdapter = this.noticeAdapter.get();
+
+            if (noticeAdapter != null && this.noticeAdapter.compareAndSet(noticeAdapter, null)) {
+                noticeAdapter.dispose();
+            }
         }).then(Mono.empty());
     }
 
@@ -282,6 +295,24 @@ public Flux<Notification> getNotifications() {
         return notifications.getEvents();
     }
 
+    @Override
+    public Flux<Notice> getNotices() {
+        NoticeAdapter notices = this.noticeAdapter.get();
+
+        if (notices == null) {
+
+            notices = new NoticeAdapter();
+
+            if (this.noticeAdapter.compareAndSet(null, notices)) {
+                notices.register(this.client);
+            } else {
+                notices = this.noticeAdapter.get();
+            }
+        }
+
+        return notices.getEvents();
+    }
+
     @Override
     public PostgresqlConnectionMetadata getMetadata() {
         return new PostgresqlConnectionMetadata(this.client.getVersion());
@@ -537,6 +568,61 @@ Flux<Notification> getEvents() {
 
     }
 
+    /**
+     * Adapter to publish {@link Notice}s.
+     */
+    static class NoticeAdapter {
+
+        private final Sinks.Many<Notice> sink = Sinks.many().multicast().directBestEffort();
+
+        @Nullable
+        private volatile Disposable subscription = null;
+
+        void dispose() {
+            Disposable subscription = this.subscription;
+            if (subscription != null && !subscription.isDisposed()) {
+                subscription.dispose();
+            }
+        }
+
+        void register(Client client) {
+
+            BaseSubscriber<NoticeResponse> subscriber = new BaseSubscriber<NoticeResponse>() {
+
+                @Override
+                protected void hookOnSubscribe(Subscription subscription) {
+                    subscription.request(Long.MAX_VALUE);
+                }
+
+                @Override
+                public void hookOnNext(NoticeResponse noticeResponse) {
+                    final Map<Field.FieldType, String> fieldsMap = new EnumMap<>(Field.FieldType.class);
+                    for (Field field : noticeResponse.getFields()) {
+                        fieldsMap.put(field.getType(), field.getValue());
+                    }
+                    NoticeAdapter.this.sink.emitNext(new Notice(fieldsMap), Sinks.EmitFailureHandler.FAIL_FAST);
+                }
+
+                @Override
+                public void hookOnError(Throwable throwable) {
+                    NoticeAdapter.this.sink.emitError(throwable, Sinks.EmitFailureHandler.FAIL_FAST);
+                }
+
+                @Override
+                public void hookOnComplete() {
+                    NoticeAdapter.this.sink.emitComplete(Sinks.EmitFailureHandler.FAIL_FAST);
+                }
+            };
+
+            this.subscription = subscriber;
+            client.addNoticeListener(subscriber);
+        }
+
+        Flux<Notice> getEvents() {
+            return this.sink.asFlux();
+        }
+    }
+
     enum EmptyTransactionDefinition implements TransactionDefinition {
 
         INSTANCE;
diff --git a/src/main/java/io/r2dbc/postgresql/api/Notice.java b/src/main/java/io/r2dbc/postgresql/api/Notice.java
@@ -0,0 +1,23 @@
+package io.r2dbc.postgresql.api;
+
+import io.r2dbc.postgresql.message.backend.Field;
+
+import java.util.Map;
+
+/**
+ * Postgres notice.
+ */
+public class Notice {
+
+    /**
+     * Notice messages by {@link Field.FieldType}.
+     */
+    public final Map<Field.FieldType, String> fields;
+
+    /**
+     * @param fields notice messages by {@link Field.FieldType}
+     */
+    public Notice(Map<Field.FieldType, String> fields) {
+        this.fields = fields;
+    }
+}
diff --git a/src/main/java/io/r2dbc/postgresql/api/PostgresqlConnection.java b/src/main/java/io/r2dbc/postgresql/api/PostgresqlConnection.java
@@ -110,6 +110,15 @@ default Mono<Long> copyIn(String sql, Publisher<ByteBuf> stdin) {
      */
     Flux<Notification> getNotifications();
 
+    /**
+     * Return a {@link Flux} of {@link Notice} received from the connection. The stream is a hot stream producing messages as they are received. Notices received by this
+     * connection are published as they are received. When the client gets {@link #close() closed}, the subscription {@link Subscriber#onComplete() completes normally}. Otherwise (transport
+     * connection disconnected unintentionally) with an {@link R2dbcNonTransientResourceException error}.
+     *
+     * @return a hot {@link Flux} of {@link Notice Notices}
+     */
+    Flux<Notice> getNotices();
+
     /**
      * Cancel currently running query by sending {@link CancelRequest} to a server.
      *
diff --git a/src/main/java/io/r2dbc/postgresql/client/Client.java b/src/main/java/io/r2dbc/postgresql/client/Client.java
@@ -18,6 +18,7 @@
 
 import io.netty.buffer.ByteBufAllocator;
 import io.r2dbc.postgresql.message.backend.BackendMessage;
+import io.r2dbc.postgresql.message.backend.NoticeResponse;
 import io.r2dbc.postgresql.message.backend.NotificationResponse;
 import io.r2dbc.postgresql.message.backend.ReadyForQuery;
 import io.r2dbc.postgresql.message.frontend.CancelRequest;
@@ -60,6 +61,27 @@ public interface Client {
      */
     void addNotificationListener(Subscriber<NotificationResponse> consumer);
 
+    /**
+     * Add a consumer of notices. Notices received by this connection are sent to the {@link Consumer notice consumer}. Note that connection errors and events such as
+     * disconnects are not visible to the {@link Consumer notice consumer}.
+     *
+     * @param consumer the consumer of notices
+     * @return a new {@link Disposable} that can be used to cancel the underlying subscription
+     * @throws IllegalArgumentException if {@code consumer} is {@code null}
+     */
+    Disposable addNoticeListener(Consumer<NoticeResponse> consumer);
+
+    /**
+     * Add a consumer of notices. Notices received by this connection are sent to the {@link Subscriber notice listener}. When the client gets {@link #close() closed}, the
+     * subscription {@link Subscriber#onComplete() completes normally}. Otherwise (transport connection disconnected unintentionally) with an {@link R2dbcNonTransientResourceException error}.
+     *
+     * @param consumer the consumer of notices
+     * @return a new {@link Disposable} that can be used to cancel the underlying subscription
+     * @throws IllegalArgumentException if {@code consumer} is {@code null}
+     * @since 1.0.0
+     */
+    void addNoticeListener(Subscriber<NoticeResponse> consumer);
+
     /**
      * Release any resources held by the {@link Client}.
      *
diff --git a/src/main/java/io/r2dbc/postgresql/client/ReactorNettyClient.java b/src/main/java/io/r2dbc/postgresql/client/ReactorNettyClient.java
@@ -112,6 +112,7 @@ public final class ReactorNettyClient implements Client {
     private final Sinks.Many<Publisher<FrontendMessage>> requestSink = Sinks.many().unicast().onBackpressureBuffer();
 
     private final Sinks.Many<NotificationResponse> notificationProcessor = Sinks.many().multicast().directBestEffort();
+    private final Sinks.Many<NoticeResponse> noticeProcessor = Sinks.many().multicast().directBestEffort();
 
     private final AtomicBoolean isClosed = new AtomicBoolean(false);
 
@@ -186,6 +187,7 @@ public Mono<Void> close() {
         return Mono.defer(() -> {
 
             this.notificationProcessor.tryEmitComplete();
+            this.noticeProcessor.tryEmitComplete();
 
             drainError(EXPECTED);
 
@@ -269,6 +271,8 @@ private boolean consumeMessage(BackendMessage message) {
 
         if (message.getClass() == NoticeResponse.class) {
 
+            this.noticeProcessor.tryEmitNext((NoticeResponse) message);
+
             this.settings.getNoticeLogLevel().log(logger, () -> this.context.getMessage(String.format("Notice: %s", toString(((NoticeResponse) message).getFields()))));
             return true;
         }
@@ -441,6 +445,16 @@ public void addNotificationListener(Subscriber<NotificationResponse> consumer) {
         this.notificationProcessor.asFlux().subscribe(consumer);
     }
 
+    @Override
+    public Disposable addNoticeListener(Consumer<NoticeResponse> consumer) {
+        return this.noticeProcessor.asFlux().subscribe(consumer);
+    }
+
+    @Override
+    public void addNoticeListener(Subscriber<NoticeResponse> consumer) {
+        this.noticeProcessor.asFlux().subscribe(consumer);
+    }
+
     @Override
     public ByteBufAllocator getByteBufAllocator() {
         return this.byteBufAllocator;
@@ -530,6 +544,7 @@ private void drainError(Supplier<? extends Throwable> supplier) {
         this.messageSubscriber.close(supplier);
 
         this.notificationProcessor.tryEmitError(supplier.get());
+        this.noticeProcessor.tryEmitError(supplier.get());
     }
 
     private final class EnsureSubscribersCompleteChannelHandler extends ChannelDuplexHandler {
diff --git a/src/test/java/io/r2dbc/postgresql/PostgresNoticeIntegrationTests.java b/src/test/java/io/r2dbc/postgresql/PostgresNoticeIntegrationTests.java
@@ -0,0 +1,91 @@
+/*
+ * 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.channel.Channel;
+import io.r2dbc.postgresql.api.Notice;
+import io.r2dbc.postgresql.api.PostgresqlConnection;
+import io.r2dbc.postgresql.api.PostgresqlResult;
+import io.r2dbc.postgresql.message.backend.Field;
+import io.r2dbc.postgresql.util.ConnectionIntrospector;
+import io.r2dbc.spi.R2dbcNonTransientResourceException;
+import java.time.Duration;
+import org.junit.jupiter.api.Test;
+import reactor.core.Disposable;
+import reactor.test.StepVerifier;
+
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.TimeUnit;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Integration tests for {@link Notice} through {@link PostgresqlConnection#getNotices()}.
+ */
+final class PostgresNoticeIntegrationTests extends AbstractIntegrationTests {
+
+    private static final String RAISE_INFO_FUNCTION =
+            "CREATE OR REPLACE FUNCTION raise_info(text)"
+                + " RETURNS void AS $$"
+                + " BEGIN"
+                + "   RAISE INFO '%', $1;"
+                + " END;"
+                + " $$ LANGUAGE plpgsql;";
+
+    @Test
+    void shouldReceivePubSubNotices() throws Exception {
+
+        BlockingQueue<Notice> notices = new LinkedBlockingQueue<>();
+
+        this.connectionFactory.create().flatMap(it -> {
+            it.getNotices().doOnNext(notices::add).subscribe();
+            return it.createStatement(RAISE_INFO_FUNCTION).execute().then()
+                    .then(it.createStatement("SELECT raise_info('Test Message')").execute().then());
+        }).block(Duration.ofSeconds(10));
+
+        Notice notice = notices.poll(10, TimeUnit.SECONDS);
+
+        assertThat(notice).isNotNull();
+        assertThat(notice.fields.containsKey(Field.FieldType.MESSAGE)).isTrue();
+        assertThat(notice.fields.get(Field.FieldType.MESSAGE)).isEqualTo("Test Message");
+    }
+
+    @Test
+    void listenShouldCompleteOnConnectionClose() {
+
+        PostgresqlConnection connection = this.connectionFactory.create().block();
+
+        connection.getNotices().as(StepVerifier::create).expectSubscription()
+            .then(() -> connection.close().subscribe())
+            .verifyComplete();
+    }
+
+    @Test
+    void listenShouldFailOnConnectionDisconnected() {
+
+        PostgresqlConnection connection = this.connectionFactory.create().block();
+
+        connection.getNotices().as(StepVerifier::create).expectSubscription()
+            .then(() -> {
+                Channel channel = ConnectionIntrospector.of(connection).getChannel();
+                channel.close();
+            })
+            .verifyError(R2dbcNonTransientResourceException.class);
+    }
+}
diff --git a/src/test/java/io/r2dbc/postgresql/api/MockPostgresqlConnection.java b/src/test/java/io/r2dbc/postgresql/api/MockPostgresqlConnection.java
@@ -77,6 +77,11 @@ public Flux<Notification> getNotifications() {
         return Flux.empty();
     }
 
+    @Override
+    public Flux<Notice> getNotices() {
+        return Flux.empty();
+    }
+
     @Override
     public Mono<Void> cancelRequest() {
         return Mono.empty();
diff --git a/src/test/java/io/r2dbc/postgresql/client/ReactorNettyClientIntegrationTests.java b/src/test/java/io/r2dbc/postgresql/client/ReactorNettyClientIntegrationTests.java
diff --git a/src/test/java/io/r2dbc/postgresql/client/TestClient.java b/src/test/java/io/r2dbc/postgresql/client/TestClient.java
