Add support for level 7 proxies (PgBouncer with pool_mode=transaction) in Postgres client.

vietj · vietj · commit aa86da09fc9b · 2023-04-03T09:56:07.000+02:00
diff --git a/vertx-pg-client/pom.xml b/vertx-pg-client/pom.xml
@@ -114,6 +114,7 @@
       <version>${jmh.version}</version>
       <scope>test</scope>
     </dependency>
+
   </dependencies>
 
   <build>
diff --git a/vertx-pg-client/src/main/asciidoc/index.adoc b/vertx-pg-client/src/main/asciidoc/index.adoc
@@ -20,6 +20,7 @@ The client is reactive and non-blocking, allowing to handle many database connec
 * SSL/TLS
 * Unix domain socket
 * HTTP/1.x CONNECT, SOCKS4a or SOCKS5 proxy support
+* Proxy (level 4 and 7) support
 
 == Usage
 
@@ -564,12 +565,32 @@ All https://www.postgresql.org/docs/current/libpq-ssl.html#LIBPQ-SSL-PROTECTION[
 
 More information can be found in the http://vertx.io/docs/vertx-core/java/#ssl[Vert.x documentation].
 
-== Using a proxy
+== Using a level 4 proxy
 
-You can also configure the client to use an HTTP/1.x CONNECT, SOCKS4a or SOCKS5 proxy.
+You can configure the client to use an HTTP/1.x CONNECT, SOCKS4a or SOCKS5 level 4 proxy.
 
 More information can be found in the http://vertx.io/docs/vertx-core/java/#_using_a_proxy_for_client_connections[Vert.x documentation].
 
+== Using a level 7 proxy
+
+Level 7 proxies can load balance queries on several connections to the actual database. When it happens, the client can be confused by the lack of session affinity and unwanted errors can happen like _ERROR: unnamed prepared statement does not exist (26000)_.
+
+Supported proxies:
+
+- https://www.pgbouncer.org[PgBouncer] configured with `_pool_mode=transaction`.
+
+You can configure the client to interact differently to the proxy:
+
+[source,$lang]
+----
+{@link examples.PgClientExamples#pgBouncer}
+----
+
+When doing so, prepared statement cannot be cached and therefore
+
+- prepared statement caching must be disabled
+- explicit prepared statement can only live within the scope of a transaction, it means you can use cursors but the prepared statement for the cursor must be created and destroyed within the scope of a transaction
+
 == Advanced pool configuration
 
 include::pool_config.adoc[]
diff --git a/vertx-pg-client/src/main/generated/io/vertx/pgclient/PgConnectOptionsConverter.java b/vertx-pg-client/src/main/generated/io/vertx/pgclient/PgConnectOptionsConverter.java
@@ -30,6 +30,11 @@ public static void fromJson(Iterable<java.util.Map.Entry<String, Object>> json,
             obj.setSslMode(io.vertx.pgclient.SslMode.valueOf((String)member.getValue()));
           }
           break;
+        case "useLayer7Proxy":
+          if (member.getValue() instanceof Boolean) {
+            obj.setUseLayer7Proxy((Boolean)member.getValue());
+          }
+          break;
       }
     }
   }
@@ -43,5 +48,6 @@ public static void toJson(PgConnectOptions obj, java.util.Map<String, Object> js
     if (obj.getSslMode() != null) {
       json.put("sslMode", obj.getSslMode().name());
     }
+    json.put("useLayer7Proxy", obj.getUseLayer7Proxy());
   }
 }
diff --git a/vertx-pg-client/src/main/java/examples/PgClientExamples.java b/vertx-pg-client/src/main/java/examples/PgClientExamples.java
@@ -732,4 +732,8 @@ public void batchReturning(SqlClient client) {
         }
       });
   }
+
+  public void pgBouncer(PgConnectOptions connectOptions) {
+    connectOptions.setUseLayer7Proxy(true);
+  }
 }
diff --git a/vertx-pg-client/src/main/java/io/vertx/pgclient/PgConnectOptions.java b/vertx-pg-client/src/main/java/io/vertx/pgclient/PgConnectOptions.java
@@ -18,6 +18,7 @@
 package io.vertx.pgclient;
 
 import io.vertx.codegen.annotations.GenIgnore;
+import io.vertx.codegen.annotations.Unstable;
 import io.vertx.core.tracing.TracingPolicy;
 import io.vertx.pgclient.impl.PgConnectionUriParser;
 import io.vertx.codegen.annotations.DataObject;
@@ -111,6 +112,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_USE_LAYER_7_PROXY = false;
   public static final Map<String, String> DEFAULT_PROPERTIES;
 
   static {
@@ -124,6 +126,7 @@ public static PgConnectOptions fromEnv() {
 
   private int pipeliningLimit = DEFAULT_PIPELINING_LIMIT;
   private SslMode sslMode = DEFAULT_SSLMODE;
+  private boolean useLayer7Proxy = DEFAULT_USE_LAYER_7_PROXY;
 
   public PgConnectOptions() {
     super();
@@ -235,6 +238,27 @@ public PgConnectOptions setSslMode(SslMode sslmode) {
     return this;
   }
 
+  /**
+   * @return whether the client interacts with a layer 7 proxy instead of a server
+   */
+  @Unstable
+  public boolean getUseLayer7Proxy() {
+    return useLayer7Proxy;
+  }
+
+  /**
+   * Set the client to use a layer 7 (application) proxy compatible protocol, set to {@code true} when the client
+   * interacts with a layer 7 proxy like PgBouncer instead of a server. Prepared statement caching must be disabled.
+   *
+   * @param useLayer7Proxy whether to use a layer 7 proxy instead of a server
+   * @return a reference to this, so the API can be used fluently
+   */
+  @Unstable
+  public PgConnectOptions setUseLayer7Proxy(boolean useLayer7Proxy) {
+    this.useLayer7Proxy = useLayer7Proxy;
+    return this;
+  }
+
   @Override
   public PgConnectOptions setSendBufferSize(int sendBufferSize) {
     return (PgConnectOptions)super.setSendBufferSize(sendBufferSize);
diff --git a/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/PgConnectionFactory.java b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/PgConnectionFactory.java
@@ -45,16 +45,21 @@ public class PgConnectionFactory extends ConnectionFactoryBase {
 
   private SslMode sslMode;
   private int pipeliningLimit;
+  private boolean useLayer7Proxy;
 
   public PgConnectionFactory(VertxInternal context, PgConnectOptions options) {
     super(context, options);
+    if (options.getCachePreparedStatements() && options.getUseLayer7Proxy()) {
+      throw new IllegalArgumentException("Prepared statement caching must be disabled when using proxy mode");
+    }
   }
 
   @Override
   protected void initializeConfiguration(SqlConnectOptions connectOptions) {
     PgConnectOptions options = (PgConnectOptions) connectOptions;
     this.pipeliningLimit = options.getPipeliningLimit();
     this.sslMode = options.isUsingDomainSocket() ? SslMode.DISABLE : options.getSslMode();
+    this.useLayer7Proxy = options.getUseLayer7Proxy();
 
     // check ssl mode here
     switch (sslMode) {
@@ -162,6 +167,6 @@ public Future<SqlConnection> connect(Context context) {
   }
 
   private PgSocketConnection newSocketConnection(EventLoopContext context, NetSocketInternal socket) {
-    return new PgSocketConnection(socket, cachePreparedStatements, preparedStatementCacheSize, preparedStatementCacheSqlFilter, pipeliningLimit, context);
+    return new PgSocketConnection(socket, cachePreparedStatements, preparedStatementCacheSize, preparedStatementCacheSqlFilter, pipeliningLimit, useLayer7Proxy, context);
   }
 }
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
@@ -50,6 +50,7 @@
 public class PgSocketConnection extends SocketConnectionBase {
 
   private PgCodec codec;
+  private final boolean useLayer7Proxy;
   public int processId;
   public int secretKey;
   public PgDatabaseMetadata dbMetaData;
@@ -59,13 +60,15 @@ public PgSocketConnection(NetSocketInternal socket,
                             int preparedStatementCacheSize,
                             Predicate<String> preparedStatementCacheSqlFilter,
                             int pipeliningLimit,
+                            boolean useLayer7Proxy,
                             EventLoopContext context) {
     super(socket, cachePreparedStatements, preparedStatementCacheSize, preparedStatementCacheSqlFilter, pipeliningLimit, context);
+    this.useLayer7Proxy = useLayer7Proxy;
   }
 
   @Override
   public void init() {
-    codec = new PgCodec();
+    codec = new PgCodec(useLayer7Proxy);
     ChannelPipeline pipeline = socket.channelHandlerContext().pipeline();
     pipeline.addBefore("handler", "codec", codec);
     super.init();
diff --git a/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/ExtendedQueryCommandCodec.java b/vertx-pg-client/src/main/java/io/vertx/pgclient/impl/codec/ExtendedQueryCommandCodec.java
@@ -48,12 +48,18 @@ void encode(PgEncoder encoder) {
           completionHandler.handle(CommandResponse.failure("Can not execute batch query with 0 sets of batch parameters."));
           return;
         } else {
+          if (encoder.useLayer7Proxy) {
+            encoder.writeParse(ps.sql, ps.bind.statement, new DataType[0]);
+          }
           for (TupleInternal param : cmd.paramsList()) {
             encoder.writeBind(ps.bind, cmd.cursorId(), param);
             encoder.writeExecute(cmd.cursorId(), cmd.fetch());
           }
         }
       } else {
+        if (encoder.useLayer7Proxy && ps.bind.statement.length == 1) {
+          encoder.writeParse(ps.sql, ps.bind.statement, new DataType[0]);
+        }
         encoder.writeBind(ps.bind, cmd.cursorId(), cmd.params());
         encoder.writeExecute(cmd.cursorId(), cmd.fetch());
       }
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
@@ -29,9 +29,9 @@ public class PgCodec extends CombinedChannelDuplexHandler<PgDecoder, PgEncoder>
 
   private final ArrayDeque<PgCommandCodec<?, ?>> inflight = new ArrayDeque<>();
 
-  public PgCodec() {
+  public PgCodec(boolean useLayer7Proxy) {
     PgDecoder decoder = new PgDecoder(inflight);
-    PgEncoder encoder = new PgEncoder(inflight);
+    PgEncoder encoder = new PgEncoder(useLayer7Proxy, inflight);
     init(decoder, encoder);
   }
 
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
@@ -62,12 +62,14 @@ final class PgEncoder extends ChannelOutboundHandlerAdapter {
   private static final byte SYNC = 'S';
 
   private final ArrayDeque<PgCommandCodec<?, ?>> inflight;
+  final boolean useLayer7Proxy;
   private ChannelHandlerContext ctx;
   private ByteBuf out;
   private final HexSequence psSeq = new HexSequence(); // used for generating named prepared statement name
   boolean closeSent;
 
-  PgEncoder(ArrayDeque<PgCommandCodec<?, ?>> inflight) {
+  PgEncoder(boolean useLayer7Proxy, ArrayDeque<PgCommandCodec<?, ?>> inflight) {
+    this.useLayer7Proxy = useLayer7Proxy;
     this.inflight = inflight;
   }
 
diff --git a/vertx-pg-client/src/test/java/io/vertx/pgclient/PgBouncerTest.java b/vertx-pg-client/src/test/java/io/vertx/pgclient/PgBouncerTest.java
@@ -0,0 +1,145 @@
+package io.vertx.pgclient;
+
+
+import io.vertx.core.*;
+import io.vertx.ext.unit.junit.VertxUnitRunner;
+import io.vertx.sqlclient.Cursor;
+import io.vertx.sqlclient.Tuple;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.testcontainers.Testcontainers;
+import org.testcontainers.containers.FixedHostPortGenericContainer;
+import org.testcontainers.containers.GenericContainer;
+import org.testcontainers.utility.DockerImageName;
+
+import java.util.*;
+import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
+
+import static org.junit.Assert.assertEquals;
+
+@RunWith(VertxUnitRunner.class)
+public class PgBouncerTest {
+
+  private FixedHostPortGenericContainer<?> pgContainer;
+  private GenericContainer<?> pgBouncerContainer;
+  private Vertx vertx;
+  private PgConnectOptions options;
+
+  @Before
+  public void setUp() {
+    pgContainer = new FixedHostPortGenericContainer<>("postgres:10.10")
+      .withFixedExposedPort(5432, 5432);
+    pgContainer.withEnv("POSTGRES_PASSWORD", "postgres");
+    pgContainer.withEnv("POSTGRES_USER", "postgres");
+    pgContainer.withEnv("POSTGRES_DB", "postgres");
+    pgContainer.withExposedPorts(5432);
+    pgContainer.withNetworkAliases("foo");
+    pgContainer.start();
+    Integer pgPort = pgContainer.getFirstMappedPort();
+    String pgHost = pgContainer.getHost();
+
+    Testcontainers.exposeHostPorts(pgPort);
+
+
+    pgBouncerContainer = new GenericContainer(DockerImageName.parse("bitnami/pgbouncer:latest"));
+    pgBouncerContainer.withEnv("POSTGRESQL_USERNAME", "postgres");
+    pgBouncerContainer.withEnv("POSTGRESQL_PASSWORD", "postgres");
+    pgBouncerContainer.withEnv("POSTGRESQL_DATABASE", "postgres");
+    pgBouncerContainer.withEnv("POSTGRESQL_HOST", "host.testcontainers.internal");
+    pgBouncerContainer.withEnv("POSTGRESQL_PORT", "" + pgPort);
+    pgBouncerContainer.withEnv("PGBOUNCER_IGNORE_STARTUP_PARAMETERS", "extra_float_digits");
+    pgBouncerContainer.withEnv("PGBOUNCER_POOL_MODE", "transaction");
+    pgBouncerContainer.withEnv("PGBOUNCER_MAX_DB_CONNECTIONS", "2");
+    pgBouncerContainer.withExposedPorts(6432);
+    pgBouncerContainer.start();
+    Integer bouncerPort = pgBouncerContainer.getFirstMappedPort();
+    String bouncerHost = pgBouncerContainer.getHost();
+
+    options = new PgConnectOptions()
+      .setHost(bouncerHost)
+      .setPort(bouncerPort)
+      .setUser("postgres")
+      .setPassword("postgres")
+      .setDatabase("postgres")
+      .addProperty("application_name", "pgbouncer-test")
+      .setPipeliningLimit(1);
+
+    vertx = Vertx.vertx();
+  }
+
+  @After
+  public void tearDown() throws Exception {
+    Testcontainers.exposeHostPorts();
+    pgBouncerContainer.stop();
+    pgContainer.stop();
+    vertx.close().toCompletionStage().toCompletableFuture().get(20, TimeUnit.SECONDS);
+  }
+
+  @Test
+  public void testPreparedQuery() throws Exception {
+    List<PgConnection> connections = new ArrayList<>();
+    int numConn = 2;
+    for (int i = 0;i < numConn;i++) {
+      connections.add(PgConnection.connect(vertx, new PgConnectOptions(options).setUseLayer7Proxy(true)).toCompletionStage().toCompletableFuture().get(20, TimeUnit.SECONDS));
+    }
+    CompositeFuture cf = CompositeFuture.join(connections.stream()
+      .map(conn -> conn.preparedQuery("select 1").execute().map(rows -> rows.iterator().next().getInteger(0)))
+      .collect(Collectors.toList()));
+    cf.toCompletionStage().toCompletableFuture().get(20, TimeUnit.SECONDS);
+    for (int i = 0;i < numConn;i++) {
+      assertEquals(1, (cf.<Object>resultAt(i)));
+    }
+  }
+
+  @Test
+  public void testPreparedBatch() throws Exception {
+    List<PgConnection> connections = new ArrayList<>();
+    int numConn = 2;
+    for (int i = 0;i < numConn;i++) {
+      connections.add(PgConnection.connect(vertx, new PgConnectOptions(options).setUseLayer7Proxy(true)).toCompletionStage().toCompletableFuture().get(20, TimeUnit.SECONDS));
+    }
+    CompositeFuture cf = CompositeFuture.join(connections.stream()
+      .map(conn -> conn.preparedQuery("select 1")
+        .executeBatch(Arrays.asList(Tuple.tuple(), Tuple.tuple()))
+        .map(rows -> rows.iterator().next().getInteger(0)))
+      .collect(Collectors.toList()));
+    cf.toCompletionStage().toCompletableFuture().get(20, TimeUnit.SECONDS);
+    for (int i = 0;i < numConn;i++) {
+      assertEquals(1, (cf.<Object>resultAt(i)));
+    }
+  }
+
+  @Test
+  public void testCursor() throws Exception {
+    List<PgConnection> connections = new ArrayList<>();
+    int numConn = 2;
+    for (int i = 0;i < numConn;i++) {
+      connections.add(PgConnection.connect(vertx, new PgConnectOptions(options).setUseLayer7Proxy(true)).toCompletionStage().toCompletableFuture().get(20, TimeUnit.SECONDS));
+    }
+    List<Future> list = new ArrayList<>();
+    for (int i = 0;i < numConn;i++) {
+      int val = i;
+      PgConnection conn = connections.get(i);
+      list.add(conn
+        .begin()
+        .compose(tx -> conn
+          .prepare("select " + val)
+          .compose(ps -> {
+            Cursor cursor = ps.cursor();
+            return cursor
+              .read(10)
+              .map(res -> res.iterator().next().getInteger(0))
+              .eventually(v -> cursor.close())
+              .eventually(v -> ps.close());
+      }).eventually(v -> tx.commit())));
+    }
+    CompositeFuture cf = CompositeFuture.join(list);
+    cf.toCompletionStage().toCompletableFuture().get(20, TimeUnit.SECONDS);
+    for (int i = 0;i < numConn;i++) {
+      assertEquals(i, (cf.<Object>resultAt(i)));
+    }
+  }
+}

Original file line number	Diff line number	Diff line change
`@@ -30,6 +30,11 @@ public static void fromJson(Iterable<java.util.Map.Entry<String, Object>> json,`
`30`	`30`	`obj.setSslMode(io.vertx.pgclient.SslMode.valueOf((String)member.getValue()));`
`31`	`31`	`}`
`32`	`32`	`break;`
	`33`	`+ case "useLayer7Proxy":`
	`34`	`+ if (member.getValue() instanceof Boolean) {`
	`35`	`+ obj.setUseLayer7Proxy((Boolean)member.getValue());`
	`36`	`+ }`
	`37`	`+ break;`
`33`	`38`	`}`
`34`	`39`	`}`
`35`	`40`	`}`
`@@ -43,5 +48,6 @@ public static void toJson(PgConnectOptions obj, java.util.Map<String, Object> js`
`43`	`48`	`if (obj.getSslMode() != null) {`
`44`	`49`	`json.put("sslMode", obj.getSslMode().name());`
`45`	`50`	`}`
	`51`	`+ json.put("useLayer7Proxy", obj.getUseLayer7Proxy());`
`46`	`52`	`}`
`47`	`53`	`}`
Original file line number	Diff line number	Diff line change
`@@ -732,4 +732,8 @@ public void batchReturning(SqlClient client) {`
`732`	`732`	`}`
`733`	`733`	`});`
`734`	`734`	`}`
	`735`	`+`
	`736`	`+ public void pgBouncer(PgConnectOptions connectOptions) {`
	`737`	`+ connectOptions.setUseLayer7Proxy(true);`
	`738`	`+ }`
`735`	`739`	`}`
Original file line number	Diff line number	Diff line change
`@@ -45,16 +45,21 @@ public class PgConnectionFactory extends ConnectionFactoryBase {`
`45`	`45`
`46`	`46`	`private SslMode sslMode;`
`47`	`47`	`private int pipeliningLimit;`
	`48`	`+ private boolean useLayer7Proxy;`
`48`	`49`
`49`	`50`	`public PgConnectionFactory(VertxInternal context, PgConnectOptions options) {`
`50`	`51`	`super(context, options);`
	`52`	`+ if (options.getCachePreparedStatements() && options.getUseLayer7Proxy()) {`
	`53`	`+ throw new IllegalArgumentException("Prepared statement caching must be disabled when using proxy mode");`
	`54`	`+ }`
`51`	`55`	`}`
`52`	`56`
`53`	`57`	`@Override`
`54`	`58`	`protected void initializeConfiguration(SqlConnectOptions connectOptions) {`
`55`	`59`	`PgConnectOptions options = (PgConnectOptions) connectOptions;`
`56`	`60`	`this.pipeliningLimit = options.getPipeliningLimit();`
`57`	`61`	`this.sslMode = options.isUsingDomainSocket() ? SslMode.DISABLE : options.getSslMode();`
	`62`	`+ this.useLayer7Proxy = options.getUseLayer7Proxy();`
`58`	`63`
`59`	`64`	`// check ssl mode here`
`60`	`65`	`switch (sslMode) {`
`@@ -162,6 +167,6 @@ public Future<SqlConnection> connect(Context context) {`
`162`	`167`	`}`
`163`	`168`
`164`	`169`	`private PgSocketConnection newSocketConnection(EventLoopContext context, NetSocketInternal socket) {`
`165`		`- return new PgSocketConnection(socket, cachePreparedStatements, preparedStatementCacheSize, preparedStatementCacheSqlFilter, pipeliningLimit, context);`
	`170`	`+ return new PgSocketConnection(socket, cachePreparedStatements, preparedStatementCacheSize, preparedStatementCacheSqlFilter, pipeliningLimit, useLayer7Proxy, context);`
`166`	`171`	`}`
`167`	`172`	`}`
Original file line number	Diff line number	Diff line change
`@@ -48,12 +48,18 @@ void encode(PgEncoder encoder) {`
`48`	`48`	`completionHandler.handle(CommandResponse.failure("Can not execute batch query with 0 sets of batch parameters."));`
`49`	`49`	`return;`
`50`	`50`	`} else {`
	`51`	`+ if (encoder.useLayer7Proxy) {`
	`52`	`+ encoder.writeParse(ps.sql, ps.bind.statement, new DataType[0]);`
	`53`	`+ }`
`51`	`54`	`for (TupleInternal param : cmd.paramsList()) {`
`52`	`55`	`encoder.writeBind(ps.bind, cmd.cursorId(), param);`
`53`	`56`	`encoder.writeExecute(cmd.cursorId(), cmd.fetch());`
`54`	`57`	`}`
`55`	`58`	`}`
`56`	`59`	`} else {`
	`60`	`+ if (encoder.useLayer7Proxy && ps.bind.statement.length == 1) {`
	`61`	`+ encoder.writeParse(ps.sql, ps.bind.statement, new DataType[0]);`
	`62`	`+ }`
`57`	`63`	`encoder.writeBind(ps.bind, cmd.cursorId(), cmd.params());`
`58`	`64`	`encoder.writeExecute(cmd.cursorId(), cmd.fetch());`
`59`	`65`	`}`
Original file line number	Diff line number	Diff line change
`@@ -29,9 +29,9 @@ public class PgCodec extends CombinedChannelDuplexHandler<PgDecoder, PgEncoder>`
`29`	`29`
`30`	`30`	`private final ArrayDeque<PgCommandCodec<?, ?>> inflight = new ArrayDeque<>();`
`31`	`31`
`32`		`- public PgCodec() {`
	`32`	`+ public PgCodec(boolean useLayer7Proxy) {`
`33`	`33`	`PgDecoder decoder = new PgDecoder(inflight);`
`34`		`- PgEncoder encoder = new PgEncoder(inflight);`
	`34`	`+ PgEncoder encoder = new PgEncoder(useLayer7Proxy, inflight);`
`35`	`35`	`init(decoder, encoder);`
`36`	`36`	`}`
`37`	`37`