Add lightweight ping syntax support

Author: mirromutth

README.md

@@ -29,7 +29,7 @@ This driver provides the following features:
 - [x] Secure connection with verification (SSL/TLS), auto-select TLS version for community and enterprise editions.
 - [x] SSL tunnel for proxy protocol of MySQL.
 - [x] Transactions with savepoint.
-- [x] Native ping command that can be verifying when argument is `ValidationDepth.REMOTE`
+- [x] Native ping can be sent via `Connection.validate(ValidationDepth.REMOTE)` and the lightweight ping syntax `/* ping */ ...`.
 - [x] Extensible, e.g. extend built-in `Codec`(s).
 - [x] MariaDB `RETURNING` clause.
 

src/main/java/io/asyncer/r2dbc/mysql/MySqlConnection.java

@@ -38,6 +38,7 @@
 import io.r2dbc.spi.ValidationDepth;
 import org.jetbrains.annotations.Nullable;
 import org.reactivestreams.Publisher;
+import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.SynchronousSink;
 
@@ -61,6 +62,8 @@ public final class MySqlConnection implements Connection, Lifecycle, ConnectionS
 
     private static final int DEFAULT_LOCK_WAIT_TIMEOUT = 50;
 
+    private static final String PING_MARKER = "/* ping */";
+
     private static final String ZONE_PREFIX_POSIX = "posix/";
 
     private static final String ZONE_PREFIX_RIGHT = "right/";
@@ -79,15 +82,28 @@ public final class MySqlConnection implements Connection, Lifecycle, ConnectionS
 
     private static final ServerVersion MARIA_10_1_1 = ServerVersion.create(10, 1, 1, true);
 
-    private static final BiConsumer<ServerMessage, SynchronousSink<Boolean>> PING = (message, sink) -> {
+    private static final Function<ServerMessage, Boolean> VALIDATE = message -> {
+        if (message instanceof CompleteMessage && ((CompleteMessage) message).isDone()) {
+            return true;
+        }
+
         if (message instanceof ErrorMessage) {
             ErrorMessage msg = (ErrorMessage) message;
             logger.debug("Remote validate failed: [{}] [{}] {}", msg.getCode(), msg.getSqlState(),
                 msg.getMessage());
-            sink.next(false);
+        } else {
+            ReferenceCountUtil.safeRelease(message);
+        }
+
+        return false;
+    };
+
+    private static final BiConsumer<ServerMessage, SynchronousSink<ServerMessage>> PING = (message, sink) -> {
+        if (message instanceof ErrorMessage) {
+            sink.next(message);
             sink.complete();
         } else if (message instanceof CompleteMessage && ((CompleteMessage) message).isDone()) {
-            sink.next(true);
+            sink.next(message);
             sink.complete();
         } else {
             ReferenceCountUtil.safeRelease(message);
@@ -190,9 +206,7 @@ public Mono<Void> beginTransaction() {
 
     @Override
     public Mono<Void> beginTransaction(TransactionDefinition definition) {
-        return Mono.defer(() -> {
-            return QueryFlow.beginTransaction(client, this, batchSupported, definition);
-        });
+        return Mono.defer(() -> QueryFlow.beginTransaction(client, this, batchSupported, definition));
     }
 
     @Override
@@ -209,9 +223,7 @@ public Mono<Void> close() {
 
     @Override
     public Mono<Void> commitTransaction() {
-        return Mono.defer(() -> {
-            return QueryFlow.doneTransaction(client, this, true, batchSupported);
-        });
+        return Mono.defer(() -> QueryFlow.doneTransaction(client, this, true, batchSupported));
     }
 
     @Override
@@ -232,6 +244,10 @@ public Mono<Void> createSavepoint(String name) {
     public MySqlStatement createStatement(String sql) {
         requireNonNull(sql, "sql must not be null");
 
+        if (sql.startsWith(PING_MARKER)) {
+            return new PingStatement(this, codecs, context);
+        }
+
         Query query = queryCache.get(sql);
 
         if (query.isSimple()) {
@@ -339,8 +355,9 @@ public Mono<Boolean> validate(ValidationDepth depth) {
                 return Mono.just(false);
             }
 
-            return client.exchange(PingMessage.INSTANCE, PING)
+            return doPingInternal()
                 .last()
+                .map(VALIDATE)
                 .onErrorResume(e -> {
                     // `last` maybe emit a NoSuchElementException, exchange maybe emit exception by Netty.
                     // But should NEVER emit any exception, so logging exception and emit false.
@@ -439,6 +456,10 @@ public Mono<Void> setStatementTimeout(Duration timeout) {
         );
     }
 
+    Flux<ServerMessage> doPingInternal() {
+        return client.exchange(PingMessage.INSTANCE, PING);
+    }
+
     boolean isSessionAutoCommit() {
         return (context.getServerStatuses() & ServerStatuses.AUTO_COMMIT) != 0;
     }

src/main/java/io/asyncer/r2dbc/mysql/MySqlStatement.java

@@ -19,6 +19,9 @@
 import io.r2dbc.spi.Statement;
 import reactor.core.publisher.Flux;
 
+import static io.asyncer.r2dbc.mysql.internal.util.AssertUtils.require;
+import static io.asyncer.r2dbc.mysql.internal.util.AssertUtils.requireNonNull;
+
 /**
  * A strongly typed implementation of {@link Statement} for the MySQL database.
  */
@@ -58,17 +61,23 @@ public interface MySqlStatement extends Statement {
      * {@inheritDoc}
      */
     @Override
-    MySqlStatement returnGeneratedValues(String... columns);
+    Flux<MySqlResult> execute();
 
     /**
      * {@inheritDoc}
      */
     @Override
-    Flux<MySqlResult> execute();
+    default MySqlStatement returnGeneratedValues(String... columns) {
+        requireNonNull(columns, "columns must not be null");
+        return this;
+    }
 
     /**
      * {@inheritDoc}
      */
     @Override
-    MySqlStatement fetchSize(int rows);
+    default MySqlStatement fetchSize(int rows) {
+        require(rows >= 0, "Fetch size must be greater or equal to zero");
+        return this;
+    }
 }

src/main/java/io/asyncer/r2dbc/mysql/MySqlStatementSupport.java

@@ -19,7 +19,6 @@
 import io.asyncer.r2dbc.mysql.internal.util.InternalArrays;
 import org.jetbrains.annotations.Nullable;
 
-import static io.asyncer.r2dbc.mysql.internal.util.AssertUtils.require;
 import static io.asyncer.r2dbc.mysql.internal.util.AssertUtils.requireNonEmpty;
 import static io.asyncer.r2dbc.mysql.internal.util.AssertUtils.requireNonNull;
 
@@ -66,12 +65,6 @@ public final MySqlStatement returnGeneratedValues(String... columns) {
         return this;
     }
 
-    @Override
-    public MySqlStatement fetchSize(int rows) {
-        require(rows >= 0, "Fetch size must be greater or equal to zero");
-        return this;
-    }
-
     @Nullable
     final String syntheticKeyName() {
         String[] columns = this.generatedColumns;

src/main/java/io/asyncer/r2dbc/mysql/PingStatement.java

@@ -0,0 +1,69 @@
+/*
+ * Copyright 2024 asyncer.io projects
+ *
+ * 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.asyncer.r2dbc.mysql;
+
+import io.asyncer.r2dbc.mysql.codec.Codecs;
+import reactor.core.publisher.Flux;
+
+/**
+ * An implementation of {@link MySqlStatement} considers the lightweight ping syntax.
+ */
+final class PingStatement implements MySqlStatement {
+
+    private final MySqlConnection connection;
+
+    private final Codecs codecs;
+
+    private final ConnectionContext context;
+
+    PingStatement(MySqlConnection connection, Codecs codecs, ConnectionContext context) {
+        this.connection = connection;
+        this.codecs = codecs;
+        this.context = context;
+    }
+
+    @Override
+    public MySqlStatement add() {
+        return this;
+    }
+
+    @Override
+    public MySqlStatement bind(int index, Object value) {
+        throw new UnsupportedOperationException("Binding parameters is not supported for ping");
+    }
+
+    @Override
+    public MySqlStatement bind(String name, Object value) {
+        throw new UnsupportedOperationException("Binding parameters is not supported for ping");
+    }
+
+    @Override
+    public MySqlStatement bindNull(int index, Class<?> type) {
+        throw new UnsupportedOperationException("Binding parameters is not supported for ping");
+    }
+
+    @Override
+    public MySqlStatement bindNull(String name, Class<?> type) {
+        throw new UnsupportedOperationException("Binding parameters is not supported for ping");
+    }
+
+    @Override
+    public Flux<MySqlResult> execute() {
+        return Flux.just(MySqlResult.toResult(false, codecs, context, null,
+            connection.doPingInternal()));
+    }
+}

src/test/java/io/asyncer/r2dbc/mysql/QueryIntegrationTestSupport.java

@@ -22,6 +22,8 @@
 import org.jetbrains.annotations.Nullable;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.condition.DisabledIf;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 import org.testcontainers.shaded.com.fasterxml.jackson.databind.JsonNode;
 import org.testcontainers.shaded.com.fasterxml.jackson.databind.ObjectMapper;
 import reactor.core.publisher.Flux;
@@ -92,6 +94,26 @@ <T> void testType(Type type, boolean valueSelect, String defined, T... values) {
                 .concatMap(value -> testOne(connection, type, valueSelect, value.orElse(null)))));
     }
 
+    @Test
+    void lightweightPing() {
+        complete(connection -> Flux.from(connection.createStatement("/* ping */ SELECT 1").execute())
+            .flatMap(result -> result.map(r -> 1))
+            .reduce(0, Integer::sum)
+            .doOnNext(number -> assertThat(number).isEqualTo(0)));
+    }
+
+    @ParameterizedTest
+    @ValueSource(strings = {
+        "/* PING */ SELECT 1", "SELECT 1 /* ping*/", " /* ping */ SELECT 1",
+        "/*ping*/ SELECT 1", "/*to ping or not to ping*/ SELECT 1"
+    })
+    void badLightweightPing(String value) {
+        complete(connection -> Flux.from(connection.createStatement(value).execute())
+            .flatMap(result -> result.map(r -> r.get(0, Integer.class)))
+            .reduce(0, Integer::sum)
+            .doOnNext(number -> assertThat(number).isEqualTo(1)));
+    }
+
     @Test
     void tinyintSigned() {
         testType(Byte.class, true, "TINYINT", (byte) 1, (byte) -1, null, Byte.MIN_VALUE, Byte.MAX_VALUE);