Preferred server type connection factory

vertx-pg-client/src/main/generated/io/vertx/pgclient/PgConnectOptionsConverter.java

@@ -25,6 +25,11 @@ public static void fromJson(Iterable<java.util.Map.Entry<String, Object>> json,
             obj.setPipeliningLimit(((Number)member.getValue()).intValue());
           }
           break;
+        case "shouldQueryServerType":
+          if (member.getValue() instanceof Boolean) {
+            obj.setShouldQueryServerType((Boolean)member.getValue());
+          }
+          break;
         case "sslMode":
           if (member.getValue() instanceof String) {
             obj.setSslMode(io.vertx.pgclient.SslMode.valueOf((String)member.getValue()));
@@ -42,6 +47,7 @@ public static void toJson(PgConnectOptions obj, JsonObject json) {
 
   public static void toJson(PgConnectOptions obj, java.util.Map<String, Object> json) {
     json.put("pipeliningLimit", obj.getPipeliningLimit());
+    json.put("shouldQueryServerType", obj.getShouldQueryServerType());
     if (obj.getSslMode() != null) {
       json.put("sslMode", obj.getSslMode().name());
     }

vertx-pg-client/src/main/java/io/vertx/pgclient/PgConnectOptions.java

@@ -111,6 +111,7 @@ public static PgConnectOptions fromEnv() {
   public static final String DEFAULT_PASSWORD = "pass";
   public static final int DEFAULT_PIPELINING_LIMIT = 256;
   public static final SslMode DEFAULT_SSLMODE = SslMode.DISABLE;
+  public static final boolean DEFAULT_SHOULD_QUERY_SERVER_TYPE = false;
   public static final Map<String, String> DEFAULT_PROPERTIES;
 
   static {
@@ -123,6 +124,7 @@ public static PgConnectOptions fromEnv() {
   }
 
   private int pipeliningLimit = DEFAULT_PIPELINING_LIMIT;
+  private boolean shouldQueryServerType = DEFAULT_SHOULD_QUERY_SERVER_TYPE;
   private SslMode sslMode = DEFAULT_SSLMODE;
 
   public PgConnectOptions() {
@@ -224,6 +226,13 @@ public SslMode getSslMode() {
     return sslMode;
   }
 
+  /**
+   * @return the value of current shouldQueryServerType
+   */
+  public boolean getShouldQueryServerType() {
+    return shouldQueryServerType;
+  }
+
   /**
    * Set {@link SslMode} for the client, this option can be used to provide different levels of secure protection.
    *
@@ -235,6 +244,19 @@ public PgConnectOptions setSslMode(SslMode sslmode) {
     return this;
   }
 
+  /**
+   * Set whether the client should query server type,
+   * If true, connection should issue implementation specific query
+   * to read {@link io.vertx.sqlclient.ServerType} of host being connected to
+   *
+   * @param shouldQueryServerType the value of shouldQueryServerType
+   * @return a reference to this, so the API can be used fluently
+   */
+  public PgConnectOptions setShouldQueryServerType(boolean shouldQueryServerType) {
+    this.shouldQueryServerType = shouldQueryServerType;
+    return this;
+  }
+
   @Override
   public PgConnectOptions setSendBufferSize(int sendBufferSize) {
     return (PgConnectOptions)super.setSendBufferSize(sendBufferSize);
@@ -464,6 +486,7 @@ protected void init() {
     this.setUser(DEFAULT_USER);
     this.setPassword(DEFAULT_PASSWORD);
     this.setDatabase(DEFAULT_DATABASE);
+    this.setShouldQueryServerType(DEFAULT_SHOULD_QUERY_SERVER_TYPE);
     this.setProperties(new HashMap<>(DEFAULT_PROPERTIES));
   }
 
@@ -493,6 +516,7 @@ public boolean equals(Object o) {
 
     if (pipeliningLimit != that.pipeliningLimit) return false;
     if (sslMode != that.sslMode) return false;
+    if (shouldQueryServerType != that.shouldQueryServerType) return false;
 
     return true;
   }
@@ -502,6 +526,7 @@ public int hashCode() {
     int result = super.hashCode();
     result = 31 * result + pipeliningLimit;
     result = 31 * result + sslMode.hashCode();
+    result = 31 * result + (shouldQueryServerType ? 1 : 0);
     return result;
   }
 

vertx-pg-client/src/main/java/io/vertx/pgclient/PgConnection.java

@@ -19,6 +19,7 @@
 
 import io.vertx.core.impl.ContextInternal;
 import io.vertx.pgclient.impl.PgConnectionImpl;
+import io.vertx.sqlclient.ServerType;
 import io.vertx.sqlclient.PreparedStatement;
 import io.vertx.sqlclient.SqlConnection;
 import io.vertx.codegen.annotations.Fluent;

vertx-pg-client/src/main/java/io/vertx/pgclient/impl/PgConnectionFactory.java

@@ -32,19 +32,27 @@
 import io.vertx.core.net.impl.NetSocketInternal;
 import io.vertx.pgclient.PgConnectOptions;
 import io.vertx.pgclient.SslMode;
+import io.vertx.sqlclient.ServerType;
 import io.vertx.sqlclient.SqlConnectOptions;
 import io.vertx.sqlclient.SqlConnection;
 import io.vertx.sqlclient.impl.Connection;
 import io.vertx.sqlclient.impl.ConnectionFactoryBase;
 import io.vertx.sqlclient.impl.tracing.QueryTracer;
 
+import java.util.Collections;
+import java.util.stream.Stream;
+
+import static io.vertx.sqlclient.ServerType.*;
+
 /**
  * @author <a href="mailto:[email protected]">Julien Viet</a>
  */
 public class PgConnectionFactory extends ConnectionFactoryBase {
 
   private SslMode sslMode;
   private int pipeliningLimit;
+  private boolean shouldQueryServerType;
+  private ServerType serverType = UNDEFINED;
 
   public PgConnectionFactory(VertxInternal context, PgConnectOptions options) {
     super(context, options);
@@ -55,6 +63,7 @@ protected void initializeConfiguration(SqlConnectOptions connectOptions) {
     PgConnectOptions options = (PgConnectOptions) connectOptions;
     this.pipeliningLimit = options.getPipeliningLimit();
     this.sslMode = options.isUsingDomainSocket() ? SslMode.DISABLE : options.getSslMode();
+    this.shouldQueryServerType = options.getShouldQueryServerType();
 
     // check ssl mode here
     switch (sslMode) {
@@ -151,16 +160,35 @@ public Future<SqlConnection> connect(Context context) {
     ContextInternal contextInternal = (ContextInternal) context;
     PromiseInternal<SqlConnection> promise = contextInternal.promise();
     connect(asEventLoopContext(contextInternal))
-      .map(conn -> {
+      .flatMap(conn -> {
         QueryTracer tracer = contextInternal.tracer() == null ? null : new QueryTracer(contextInternal.tracer(), options);
         PgConnectionImpl pgConn = new PgConnectionImpl(this, contextInternal, conn, tracer, null);
         conn.init(pgConn);
-        return (SqlConnection)pgConn;
+        serverType = ((PgSocketConnection) conn).serverType;
+        if (serverType == UNDEFINED && shouldQueryServerType) {
+          String paramStatus = conn.getDatabaseMetaData().majorVersion() >= 14 ? "in_hot_standby" : "transaction_read_only";
+          return pgConn.query(String.format("SHOW %s;", paramStatus)).execute().map(pgRowSet -> {
+            pgRowSet.forEach(row ->
+              serverType = "off".equalsIgnoreCase(row.getString(paramStatus)) ? PRIMARY : REPLICA
+            );
+            return pgConn;
+          });
+        } else {
+          return Future.succeededFuture((SqlConnection) pgConn);
+        }
       })
       .onComplete(promise);
     return promise.future();
   }
 
+  /**
+   * {@inheritDoc}
+   */
+  @Override
+  public ServerType getServerType() {
+    return serverType;
+  }
+
   private PgSocketConnection newSocketConnection(EventLoopContext context, NetSocketInternal socket) {
     return new PgSocketConnection(socket, cachePreparedStatements, preparedStatementCacheSize, preparedStatementCacheSqlFilter, pipeliningLimit, context);
   }

vertx-pg-client/src/main/java/io/vertx/pgclient/impl/PgConnectionImpl.java

@@ -21,6 +21,7 @@
 import io.vertx.pgclient.PgConnectOptions;
 import io.vertx.pgclient.PgConnection;
 import io.vertx.pgclient.PgNotification;
+import io.vertx.sqlclient.ServerType;
 import io.vertx.sqlclient.impl.Connection;
 import io.vertx.sqlclient.impl.Notification;
 import io.vertx.sqlclient.impl.SqlConnectionImpl;

vertx-pg-client/src/main/java/io/vertx/pgclient/impl/PgPoolImpl.java

@@ -28,6 +28,7 @@
 import io.vertx.pgclient.*;
 import io.vertx.pgclient.spi.PgDriver;
 import io.vertx.sqlclient.PoolOptions;
+import io.vertx.sqlclient.ServerRequirement;
 import io.vertx.sqlclient.SqlConnectOptions;
 import io.vertx.sqlclient.SqlConnection;
 import io.vertx.sqlclient.impl.Connection;
@@ -71,8 +72,12 @@ public static PgPoolImpl create(final VertxInternal vertx, boolean pipelined, Li
     int pipeliningLimit = pipelined ? baseConnectOptions.getPipeliningLimit() : 1;
     PgPoolImpl pool = new PgPoolImpl(vx, baseConnectOptions, tracer, metrics, pipeliningLimit, poolOptions);
     PgDriver driver = new PgDriver();
-    List<ConnectionFactory> lst = servers.stream().map(options -> driver.createConnectionFactory(vx, options)).collect(Collectors.toList());
-    ConnectionFactory factory = ConnectionFactory.roundRobinSelector(lst);
+    List<ConnectionFactory> lst = servers.stream()
+      .map(options -> driver.createConnectionFactory(vx, PgConnectOptions.wrap(options)
+        .setShouldQueryServerType(poolOptions.getServerRequirement() != ServerRequirement.ANY))
+      )
+      .collect(Collectors.toList());
+    ConnectionFactory factory = ConnectionFactory.selector(lst, poolOptions.getServerRequirement());
     pool.connectionProvider(factory::connect);
     pool.init();
     CloseFuture closeFuture = pool.closeFuture();

vertx-pg-client/src/main/java/io/vertx/pgclient/impl/PgSocketConnection.java

@@ -28,6 +28,7 @@
 import io.vertx.core.net.impl.NetSocketInternal;
 import io.vertx.pgclient.PgException;
 import io.vertx.pgclient.impl.codec.PgCodec;
+import io.vertx.sqlclient.ServerType;
 import io.vertx.pgclient.impl.codec.TxFailedEvent;
 import io.vertx.sqlclient.impl.*;
 import io.vertx.sqlclient.impl.command.*;
@@ -36,6 +37,8 @@
 import java.util.Map;
 import java.util.function.Predicate;
 
+import static io.vertx.sqlclient.ServerType.UNDEFINED;
+
 /**
  * @author <a href="mailto:[email protected]">Julien Viet</a>
  */
@@ -45,6 +48,8 @@ public class PgSocketConnection extends SocketConnectionBase {
   public int processId;
   public int secretKey;
   public PgDatabaseMetadata dbMetaData;
+  // TODO: consider defining getter on SocketConnectionBase
+  public ServerType serverType = UNDEFINED;
 
   public PgSocketConnection(NetSocketInternal socket,
                             boolean cachePreparedStatements,

vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/InitCommandCodec.java

@@ -27,6 +27,9 @@
 import io.vertx.sqlclient.impl.command.CommandResponse;
 import io.vertx.sqlclient.impl.command.InitCommand;
 
+import static io.vertx.sqlclient.ServerType.PRIMARY;
+import static io.vertx.sqlclient.ServerType.REPLICA;
+
 class InitCommandCodec extends PgCommandCodec<Connection, InitCommand> {
 
   private PgEncoder encoder;
@@ -87,6 +90,9 @@ public void handleParameterStatus(String key, String value) {
     if(key.equals("server_version")) {
       ((PgSocketConnection)cmd.connection()).dbMetaData = new PgDatabaseMetadata(value);
     }
+    if(key.equals("in_hot_standby")) {
+      ((PgSocketConnection)cmd.connection()).serverType = "off".equalsIgnoreCase(value) ? PRIMARY : REPLICA;
+    }
   }
 
   @Override