io.r2dbc.spi.Statement Java Examples

@Test
default void returnGeneratedValues() {

    getJdbcOperations().execute("DROP TABLE test");
    getJdbcOperations().execute(getCreateTableWithAutogeneratedKey());

    Mono.from(getConnectionFactory().create())
        .flatMapMany(connection -> {
            Statement statement = connection.createStatement(getInsertIntoWithAutogeneratedKey());

            statement.returnGeneratedValues();

            return Flux.from(statement
                .execute())
                .concatWith(close(connection)).flatMap(it -> it.map((row, rowMetadata) -> row.get(0)));
        })
        .as(StepVerifier::create)
        .expectNextCount(1)
        .verifyComplete();
}
 

@Override
@Test
public void prepareStatement() {
  Mono.from(getConnectionFactory().create())
      .delayUntil(c -> c.beginTransaction())
      .flatMapMany(connection -> {
        Statement statement = connection.createStatement(
            String.format("INSERT INTO test (value) VALUES(%s)", getPlaceholder(0)));

        IntStream.range(0, 10)
            .forEach(i -> statement.bind(getIdentifier(0), i).add());

        return Flux.from(statement
            .execute())
            .concatWith(close(connection));
      })
      .as(StepVerifier::create)
      .expectNextCount(10).as("values from insertions")
      .verifyComplete();
}
 

@Test
default void bindFails() {
    Mono.from(getConnectionFactory().create())
        .flatMap(connection -> {

            Statement statement = connection.createStatement(String.format("INSERT INTO test VALUES(%s)", getPlaceholder(0)));
            assertThrows(IllegalArgumentException.class, () -> statement.bind(0, null), "bind(0, null) should fail");
            assertThrows(IndexOutOfBoundsException.class, () -> statement.bind(99, ""), "bind(nonexistent-index, null) should fail");
            assertThrows(IllegalArgumentException.class, () -> bind(statement, getIdentifier(0), null), "bind(identifier, null) should fail");
            assertThrows(IllegalArgumentException.class, () -> bind(statement, getIdentifier(0), Class.class), "bind(identifier, Class.class) should fail");
            assertThrows(IllegalArgumentException.class, () -> statement.bind("unknown", ""), "bind(unknown-placeholder, \"\") should fail");
            return close(connection);
        })
        .as(StepVerifier::create)
        .verifyComplete();
}
 

@Test
default void bindNull() {
    Mono.from(getConnectionFactory().create())
        .flatMapMany(connection -> {
            Statement statement = connection
                .createStatement(String.format("INSERT INTO test VALUES(%s)", getPlaceholder(0)));
            bindNull(statement, getIdentifier(0), Integer.class);
            return Flux.from(statement.add()
                .execute())
                .flatMap(TestKit::extractRowsUpdated)
                .concatWith(close(connection));
        })
        .as(StepVerifier::create)
        .expectNextCount(1).as("rows inserted")
        .verifyComplete();
}
 

@Test
default void prepareStatement() {
    Mono.from(getConnectionFactory().create())
        .flatMapMany(connection -> {
            Statement statement = connection.createStatement(String.format("INSERT INTO test VALUES(%s)", getPlaceholder(0)));

            IntStream.range(0, 10)
                .forEach(i -> bind(statement, getIdentifier(0), i).add());

            return Flux.from(statement
                .execute())
                .flatMap(TestKit::extractRowsUpdated)
                .concatWith(close(connection));
        })
        .as(StepVerifier::create)
        .expectNextCount(10).as("values from insertions")
        .verifyComplete();
}
 

@Test
default void bindNullFails() {
    Mono.from(getConnectionFactory().create())
        .flatMap(connection -> {

            Statement statement = connection.createStatement(String.format("INSERT INTO test VALUES(%s)", getPlaceholder(0)));
            assertThrows(IllegalArgumentException.class, () -> statement.bindNull(null, String.class), "bindNull(null, …) should fail");
            assertThrows(IllegalArgumentException.class, () -> bind(statement, getIdentifier(0), null), "bindNull(identifier, null) should fail");
            return close(connection);
        })
        .as(StepVerifier::create)
        .verifyComplete();
}
 

static Statement bind(Statement statement, Object identifier, Object value) {
    Assert.requireNonNull(identifier, "Identifier must not be null");
    if (identifier instanceof String) {
        return statement.bind((String) identifier, value);
    }
    if (identifier instanceof Integer) {
        return statement.bind((Integer) identifier, value);
    }
    throw new IllegalArgumentException(String.format("Identifier %s must be a String or Integer. Was: %s", identifier, identifier.getClass().getName()));
}
 

@Test
default void transactionRollback() {
    getJdbcOperations().execute("INSERT INTO test VALUES (100)");

    Mono.from(getConnectionFactory().create())
        .flatMapMany(connection -> Mono.from(connection

            .beginTransaction())
            .<Object>thenMany(Flux.from(connection.createStatement("SELECT value FROM test")
                .execute())
                .flatMap(TestKit::extractColumns))

            .concatWith(Flux.defer(() -> {
                Statement statement = connection.createStatement(String.format("INSERT INTO test VALUES (%s)", getPlaceholder(0)));
                bind(statement, getIdentifier(0), 200);
                return statement.execute();
            })
                .flatMap(TestKit::extractRowsUpdated))
            .concatWith(Flux.from(connection.createStatement("SELECT value FROM test")
                .execute())
                .flatMap(TestKit::extractColumns))

            .concatWith(connection.rollbackTransaction())
            .concatWith(Flux.from(connection.createStatement("SELECT value FROM test")
                .execute())
                .flatMap(TestKit::extractColumns))

            .concatWith(close(connection)))
        .as(StepVerifier::create)
        .expectNext(Collections.singletonList(100)).as("value from select")
        .expectNext(1).as("rows inserted")
        .expectNext(Arrays.asList(100, 200)).as("values from select")
        .expectNext(Collections.singletonList(100)).as("value from select")
        .verifyComplete();
}
 

@Test
default void transactionCommit() {
    getJdbcOperations().execute("INSERT INTO test VALUES (100)");

    Mono.from(getConnectionFactory().create())
        .flatMapMany(connection -> Mono.from(connection

            .beginTransaction())
            .<Object>thenMany(Flux.from(connection.createStatement("SELECT value FROM test")
                .execute())
                .flatMap(TestKit::extractColumns))

            .concatWith(Flux.defer(() -> {
                Statement statement = connection.createStatement(String.format("INSERT INTO test VALUES (%s)", getPlaceholder(0)));
                bind(statement, getIdentifier(0), 200);
                return statement.execute();
            })
                .flatMap(TestKit::extractRowsUpdated))
            .concatWith(Flux.from(connection.createStatement("SELECT value FROM test")
                .execute())
                .flatMap(TestKit::extractColumns))

            .concatWith(connection.commitTransaction())
            .concatWith(Flux.from(connection.createStatement("SELECT value FROM test")
                .execute())
                .flatMap(TestKit::extractColumns))

            .concatWith(close(connection)))
        .as(StepVerifier::create)
        .expectNext(Collections.singletonList(100)).as("value from select")
        .expectNext(1).as("rows inserted")
        .expectNext(Arrays.asList(100, 200)).as("values from select")
        .expectNext(Arrays.asList(100, 200)).as("values from select")
        .verifyComplete();
}
 

@Test
default void returnGeneratedValuesFails() {

    Mono.from(getConnectionFactory().create())
        .flatMapMany(connection -> {
            Statement statement = connection.createStatement("INSERT INTO test");

            assertThrows(IllegalArgumentException.class, () -> statement.returnGeneratedValues((String[]) null));
            return close(connection);
        })
        .as(StepVerifier::create)
        .verifyComplete();
}
 

@Test
default void prepareStatementWithIncompleteBindingFails() {
    Mono.from(getConnectionFactory().create())
        .flatMapMany(connection -> {
            Statement statement = connection.createStatement(String.format("INSERT INTO test VALUES(%s,%s)", getPlaceholder(0), getPlaceholder(1)));

            bind(statement, getIdentifier(0), 0);

            assertThrows(IllegalStateException.class, statement::execute);
            return close(connection);
        })
        .as(StepVerifier::create)
        .verifyComplete();
}
 

@Test
default void prepareStatementWithIncompleteBatchFails() {
    Mono.from(getConnectionFactory().create())
        .flatMapMany(connection -> {
            Statement statement = connection.createStatement(String.format("INSERT INTO test VALUES(%s,%s)", getPlaceholder(0), getPlaceholder(1)));

            bind(statement, getIdentifier(0), 0);

            assertThrows(IllegalStateException.class, statement::add);
            return close(connection);
        })
        .as(StepVerifier::create)
        .verifyComplete();
}
 

@Test
default void clobInsert() {
    Mono.from(getConnectionFactory().create())
        .flatMapMany(connection -> {
            Statement statement = connection.createStatement(String.format("INSERT INTO clob_test VALUES (%s)", getPlaceholder(0)));
            bind(statement, getIdentifier(0), Clob.from(Mono.just("test-value")));
            return Flux.from(statement.execute())
                .flatMap(Result::getRowsUpdated)
                .concatWith(close(connection));
        })
        .as(StepVerifier::create)
        .expectNextCount(1).as("rows inserted")
        .verifyComplete();
}
 

@Test
default void blobInsert() {
    Mono.from(getConnectionFactory().create())
        .flatMapMany(connection -> {
            Statement statement = connection.createStatement(String.format("INSERT INTO blob_test VALUES (%s)", getPlaceholder(0)));
            bind(statement, getIdentifier(0), Blob.from(Mono.just(StandardCharsets.UTF_8.encode("test-value"))));
            return Flux.from(statement.execute())
                .flatMap(TestKit::extractRowsUpdated)
                .concatWith(close(connection));
        })
        .as(StepVerifier::create)
        .expectNextCount(1).as("rows inserted")
        .verifyComplete();
}
 

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

@Test
public void executeStatementReturnsWorkingStatementWithCorrectQuery() {
  SpannerConnection connection
      = new SpannerConnection(this.mockClient, TEST_SESSION, TEST_CONFIG);
  String sql = "select book from library";
  PartialResultSet partialResultSet = PartialResultSet.newBuilder()
      .setMetadata(ResultSetMetadata.newBuilder().setRowType(StructType.newBuilder()
          .addFields(
              Field.newBuilder().setName("book")
                  .setType(Type.newBuilder().setCode(TypeCode.STRING)))))
      .addValues(Value.newBuilder().setStringValue("Odyssey"))
      .build();

  when(this.mockClient.executeStreamingSql(
        any(StatementExecutionContext.class), eq(sql), eq(EMPTY_STRUCT), eq(EMPTY_TYPE_MAP)))
      .thenReturn(Flux.just(makeBookPrs("Odyssey")));

  Statement statement = connection.createStatement(sql);
  assertThat(statement).isInstanceOf(SpannerStatement.class);

  StepVerifier.create(
      ((Flux<SpannerResult>) statement.execute())
          .flatMap(res -> res.map((r, m) -> (String) r.get(0))))
      .expectNext("Odyssey")
      .expectComplete()
      .verify();

  verify(this.mockClient).executeStreamingSql(any(StatementExecutionContext.class), eq(sql),
      eq(EMPTY_STRUCT), eq(EMPTY_TYPE_MAP));

  // Single use READ query doesn't need these round trips below.
  verify(this.mockClient, times(0)).beginTransaction(eq(TEST_SESSION_NAME), any());
  verify(this.mockClient, times(0)).commitTransaction(eq(TEST_SESSION_NAME), any());
}
 

@Override
public Publisher<? extends Result> execute() {
  if (this.statementType == StatementType.DDL) {
    return this.client
        .executeDdl(
            this.config.getFullyQualifiedDatabaseName(),
            Collections.singletonList(this.sql),
            this.config.getDdlOperationTimeout(),
            this.config.getDdlOperationPollInterval())
        .map(operation -> new SpannerResult(Flux.empty(), Mono.just(0)));
  } else if (this.statementType == StatementType.DML && !this.ctx.isTransactionPartitionedDml()) {

    List<ExecuteBatchDmlRequest.Statement> dmlStatements =
        this.statementBindings.getBindings().stream()
            .map(struct ->
                ExecuteBatchDmlRequest.Statement.newBuilder()
                    .setSql(this.sql)
                    .setParams(struct)
                    .putAllParamTypes(this.statementBindings.getTypes())
                    .build())
            .collect(Collectors.toList());

    return this.client.executeBatchDml(this.ctx, dmlStatements)
        .map(partialResultSet -> Math.toIntExact(partialResultSet.getStats().getRowCountExact()))
        .map(rowCount -> new SpannerResult(Flux.empty(), Mono.just(rowCount)));
  }

  Flux<Struct> structFlux = Flux.fromIterable(this.statementBindings.getBindings());
  return structFlux.flatMap(this::runStreamingSql);
}
 

static Statement bindNull(Statement statement, Object identifier, Class<?> type) {
    Assert.requireNonNull(identifier, "Identifier must not be null");
    if (identifier instanceof String) {
        return statement.bindNull((String) identifier, type);
    }
    if (identifier instanceof Integer) {
        return statement.bindNull((Integer) identifier, type);
    }
    throw new IllegalArgumentException(String.format("Identifier %s must be a String or Integer. Was: %s", identifier, identifier.getClass().getName()));
}
 

@Override
public Statement bind(String name, Object value) {
    Assert.requireNonNull(name, "name must not be null");
    Assert.requireNonNull(value, "value must not be null");

    getCurrent().put(name, value);
    return this;
}
 

@Override
public Statement bindNull(String name, Class<?> type) {
    Assert.requireNonNull(name, "name must not be null");
    Assert.requireNonNull(type, "value must not be null");

    getCurrent().put(name, type);
    return this;
}
 

@Override
public Statement returnGeneratedValues(String... columns) {
    this.generatedValuesColumns = columns;
    return this;
}
 

@Override
public Statement add() {
  this.statementBindings.completeBinding();
  return this;
}
 

@Override
public Statement bindNull(int i, Class<?> type) {
  throw new UnsupportedOperationException("Only named parameters are supported");
}
 

@Override
public Statement bindNull(String identifier, Class<?> type) {
  return bind(identifier, new TypedNull(type));
}
 

@Override
public Statement bind(int i, Object o) {
  throw new UnsupportedOperationException("Only named parameters are supported");
}
 

@Override
public Statement bind(String identifier, Object value) {
  this.statementBindings.createBind(identifier, value);
  return this;
}