io.r2dbc.spi.Connection Java Examples

@Test
public void reusesUntilConnectionFactoryIsClosed() {
    String url = "r2dbc:tc:postgresql:///db?TC_IMAGE_TAG=10-alpine";
    ConnectionFactory connectionFactory = ConnectionFactories.get(url);

    Integer updated = Flux
        .usingWhen(
            connectionFactory.create(),
            connection -> {
                return Mono
                    .from(connection.createStatement("CREATE TABLE test(id integer PRIMARY KEY)").execute())
                    .thenMany(connection.createStatement("INSERT INTO test(id) VALUES(123)").execute())
                    .flatMap(Result::getRowsUpdated);
            },
            Connection::close
        )
        .blockFirst();

    assertThat(updated).isEqualTo(1);

    Flux<Long> select = Flux
        .usingWhen(
            Flux.defer(connectionFactory::create),
            connection -> {
                return Flux
                    .from(connection.createStatement("SELECT COUNT(*) FROM test").execute())
                    .flatMap(it -> it.map((row, meta) -> (Long) row.get(0)));
            },
            Connection::close
        );

    Long rows = select.blockFirst();

    assertThat(rows).isEqualTo(1);

    close(connectionFactory);

    Assertions
        .assertThatThrownBy(select::blockFirst)
        .isInstanceOf(PostgresqlException.class)
        // relation "X" does not exists
        // https://github.com/postgres/postgres/blob/REL_10_0/src/backend/utils/errcodes.txt#L349
        .returns("42P01", e -> ((PostgresqlException) e).getErrorDetails().getCode());
}
 

private static Mono<Void> testTimeDuration(Connection connection, Duration origin, LocalTime time) {
    return Mono.from(connection.createStatement("INSERT INTO test VALUES(DEFAULT,?)")
        .bind(0, origin)
        .returnGeneratedValues("id")
        .execute())
        .flatMapMany(QueryIntegrationTestSupport::extractId)
        .concatMap(id -> connection.createStatement("SELECT value FROM test WHERE id=?")
            .bind(0, id)
            .execute())
        .<Optional<LocalTime>>flatMap(r -> extractOptionalField(r, LocalTime.class))
        .map(Optional::get)
        .doOnNext(t -> assertThat(t).isEqualTo(time))
        .then(Mono.from(connection.createStatement("DELETE FROM test WHERE id>0")
            .execute()))
        .flatMap(QueryIntegrationTestSupport::extractRowsUpdated)
        .then();
}
 

@Test
public void testSessionManagement() {
  assertThat(this.connectionFactory).isInstanceOf(SpannerConnectionFactory.class);

  Mono<Connection> connection = (Mono<Connection>) this.connectionFactory.create();
  SpannerConnection spannerConnection = (SpannerConnection) connection.block();
  String activeSessionName = spannerConnection.getSessionName();

  List<String> activeSessions = getSessionNames();
  assertThat(activeSessions).contains(activeSessionName);

  Mono.from(spannerConnection.close()).block();

  activeSessions = getSessionNames();
  assertThat(activeSessions).doesNotContain(activeSessionName);
}
 

protected void runTestQuery(ConnectionFactory connectionFactory) {
    try {
        int expected = 42;
        Number result = Flux
            .usingWhen(
                connectionFactory.create(),
                connection -> connection.createStatement(createTestQuery(expected)).execute(),
                Connection::close
            )
            .flatMap(it -> it.map((row, meta) -> (Number) row.get(0)))
            .blockFirst();

        assertThat(result)
            .isNotNull()
            .returns(expected, Number::intValue);
    } finally {
        if (connectionFactory instanceof Closeable) {
            Mono.from(((Closeable) connectionFactory).close()).block();
        }
    }
}
 

/**
 * Initializes the integration test environment for the Spanner R2DBC dialect.
 */
@BeforeEach
public void initializeTestEnvironment() {
  Connection connection = Mono.from(connectionFactory.create()).block();

  this.databaseClient = DatabaseClient.create(connectionFactory);

  if (SpannerTestUtils.tableExists(connection, "PRESIDENT")) {
    this.databaseClient.execute("DROP TABLE PRESIDENT")
        .fetch()
        .rowsUpdated()
        .block();
  }

  this.databaseClient.execute(
      "CREATE TABLE PRESIDENT ("
          + "  NAME STRING(256) NOT NULL,"
          + "  START_YEAR INT64 NOT NULL"
          + ") PRIMARY KEY (NAME)")
      .fetch()
      .rowsUpdated()
      .block();
}
 

private static Mono<Void> testTimeDuration(Connection connection, Duration origin, LocalTime time) {
    return Mono.from(connection.createStatement("INSERT INTO test VALUES(DEFAULT,?)")
        .bind(0, origin)
        .returnGeneratedValues("id")
        .execute())
        .flatMapMany(QueryIntegrationTestSupport::extractId)
        .concatMap(id -> connection.createStatement("SELECT value FROM test WHERE id=?")
            .bind(0, id)
            .execute())
        .<Optional<LocalTime>>flatMap(r -> extractOptionalField(r, LocalTime.class))
        .map(Optional::get)
        .doOnNext(t -> assertThat(t).isEqualTo(time))
        .then(Mono.from(connection.createStatement("DELETE FROM test WHERE id>0")
            .execute()))
        .flatMap(QueryIntegrationTestSupport::extractRowsUpdated)
        .then();
}
 

@Override
public Publisher<? extends Connection> create() {
    return new ConnectionPublisher(
        () -> {
            if (future == null) {
                synchronized (this) {
                    if (future == null) {
                        future = CompletableFuture.supplyAsync(() -> {
                            R2DBCDatabaseContainer container = containerProvider.createContainer(options);
                            container.start();
                            return container;
                        }, EXECUTOR);
                    }
                }
            }
            return future.thenApply(it -> {
                return ConnectionFactories.find(
                    it.configure(options)
                );
            });
        }
    );
}
 

private static void executeDml(Function<Connection, Statement> statementFunc) {
  Mono.from(connectionFactory.create())
      .delayUntil(c -> c.beginTransaction())
      .delayUntil(c -> Flux.from(statementFunc.apply(c).execute())
          .flatMapSequential(r -> Mono.from(r.getRowsUpdated())))
      .delayUntil(c -> c.commitTransaction())
      .delayUntil(c -> c.close())
      .block();
}
 

@Override
void enter() {
    this.future = futureSupplier.get();

    future.whenComplete((connectionFactory, e) -> {
        if (e != null) {
            actual.onSubscribe(EmptySubscription.INSTANCE);
            actual.onError(e);
            return;
        }

        Publisher<? extends Connection> publisher = connectionFactory.create();
        transitionTo(new ProxySubscriptionState(publisher));
    });
}
 

@Override
public Mono<Connection> create() {
    if (this.connection == null) {
        throw new AssertionError("Unexpected call to create()");
    }

    return this.connection;
}
 

/**
 * Returns true if the Spanner table exists; false if not.
 */
public static boolean tableExists(Connection connection, String tableName) {
  return Mono.from(connection.createStatement(
      "SELECT table_name FROM information_schema.tables WHERE table_name = @name")
      .bind("name", tableName)
      .execute())
      .flatMapMany(result -> result.map((r, m) -> r))
      .hasElements()
      .block();
}
 

/**
 * Executes a DML query and returns the rows updated.
 */
static int executeDmlQuery(ConnectionFactory connectionFactory, String sql) {
  Connection connection = Mono.from(connectionFactory.create()).block();

  Mono.from(connection.beginTransaction()).block();
  int rowsUpdated = Mono.from(connection.createStatement(sql).execute())
      .flatMap(result -> Mono.from(result.getRowsUpdated()))
      .block();
  Mono.from(connection.commitTransaction()).block();

  return rowsUpdated;
}
 

private MockConnectionFactory(@Nullable Mono<Connection> connection) {
    this.connection = connection;
}
 

static <T> Mono<T> close(Connection connection) {
    return Mono.from(connection
        .close())
        .then(Mono.empty());
}
 

@Override
public void subscribe(Subscriber<? super Connection> actual) {
    actual.onSubscribe(new StateMachineSubscription(actual));
}
 

StateMachineSubscription(Subscriber<? super Connection> actual) {
    this.actual = actual;
    subscriptionState = new WaitRequestSubscriptionState();
}
 

ProxySubscriptionState(Publisher<? extends Connection> publisher) {
    this.publisher = publisher;
}
 

@Override
public void onNext(Connection connection) {
    actual.onNext(connection);
}
 

private static <T> Mono<T> close(Connection connection) {
  return Mono.from(connection
      .close())
      .then(Mono.empty());
}
 

private <T> Mono<T> close(Connection connection) {
    return Mono.from(connection.close())
      .then(Mono.empty());
}