PgBouncerTest.java

package io.vertx.tests.pgclient;


import io.vertx.core.CompositeFuture;
import io.vertx.core.Future;
import io.vertx.core.Vertx;
import io.vertx.ext.unit.junit.VertxUnitRunner;
import io.vertx.pgclient.PgConnectOptions;
import io.vertx.pgclient.PgConnection;
import io.vertx.sqlclient.Cursor;
import io.vertx.sqlclient.Tuple;
import org.apache.commons.compress.utils.IOUtils;
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.images.builder.ImageFromDockerfile;
import org.testcontainers.images.builder.Transferable;

import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

import static org.junit.Assert.assertEquals;
import static org.testcontainers.containers.BindMode.READ_ONLY;

@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();

    Testcontainers.exposeHostPorts(pgPort);

    pgBouncerContainer = new GenericContainer<>(
      new ImageFromDockerfile()
        .withFileFromTransferable("Dockerfile", resourceTransferable("pgBouncer/Dockerfile")))
      .withClasspathResourceMapping("pgBouncer/pgbouncer.ini", "/etc/pgbouncer/pgbouncer.ini", READ_ONLY)
      .withClasspathResourceMapping("pgBouncer/userlist.txt", "/etc/pgbouncer/userlist.txt", READ_ONLY)
      .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();
  }

  private static Transferable resourceTransferable(String resourceName) {
    try (InputStream stream = PgBouncerTest.class.getClassLoader().getResourceAsStream(resourceName)) {
      assert stream != null;
      return Transferable.of(IOUtils.toByteArray(stream));
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
  }

  @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 = Future.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 = Future.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(() -> cursor.close())
              .eventually(() -> ps.close());
      }).eventually(() -> tx.commit())));
    }
    CompositeFuture cf = Future.join(list);
    cf.toCompletionStage().toCompletableFuture().get(20, TimeUnit.SECONDS);
    for (int i = 0;i < numConn;i++) {
      assertEquals(i, (cf.<Object>resultAt(i)));
    }
  }
}