com.google.cloud.spanner.ReadOnlyTransaction Java Examples

@Test
public void readOnlyTransactionPartitionedDmlTest() {

	this.expectedException.expectMessage("A read-only transaction template cannot execute partitioned" +
			" DML.");

	ReadOnlyTransaction readOnlyTransaction = mock(ReadOnlyTransaction.class);
	when(this.databaseClient.readOnlyTransaction(
			eq(TimestampBound.ofReadTimestamp(Timestamp.ofTimeMicroseconds(333)))))
					.thenReturn(readOnlyTransaction);

	this.spannerTemplate
			.performReadOnlyTransaction((spannerOperations) -> {
				spannerOperations.executePartitionedDmlStatement(Statement.of("fail"));
				return null;
			}, new SpannerReadOptions()
					.setTimestamp(Timestamp.ofTimeMicroseconds(333)));
}
 

private void preparePkMetadata(ReadOnlyTransaction tx, List<Struct> rows) {
  Type type =
      Type.struct(
          Type.StructField.of("table_name", Type.string()),
          Type.StructField.of("column_name", Type.string()),
          Type.StructField.of("column_ordering", Type.string()));
  when(tx.executeQuery(
          argThat(
              new ArgumentMatcher<Statement>() {

                @Override
                public boolean matches(Statement argument) {
                  if (!(argument instanceof Statement)) {
                    return false;
                  }
                  Statement st = (Statement) argument;
                  return st.getSql().contains("information_schema.index_columns");
                }
              })))
      .thenReturn(ResultSets.forRows(type, rows));
}
 

private void prepareColumnMetadata(ReadOnlyTransaction tx, List<Struct> rows) {
  Type type =
      Type.struct(
          Type.StructField.of("table_name", Type.string()),
          Type.StructField.of("column_name", Type.string()),
          Type.StructField.of("spanner_type", Type.string()),
          Type.StructField.of("cells_mutated", Type.int64()));
  when(tx.executeQuery(
          argThat(
              new ArgumentMatcher<Statement>() {

                @Override
                public boolean matches(Statement argument) {
                  if (!(argument instanceof Statement)) {
                    return false;
                  }
                  Statement st = (Statement) argument;
                  return st.getSql().contains("information_schema.columns");
                }
              })))
      .thenReturn(ResultSets.forRows(type, rows));
}
 

@Before
@SuppressWarnings("unchecked")
public void setUp() throws Exception {
  MockitoAnnotations.initMocks(this);
  serviceFactory = new FakeServiceFactory();

  ReadOnlyTransaction tx = mock(ReadOnlyTransaction.class);
  when(serviceFactory.mockDatabaseClient().readOnlyTransaction()).thenReturn(tx);

  // Capture batches sent to writeAtLeastOnce.
  when(serviceFactory.mockDatabaseClient().writeAtLeastOnce(mutationBatchesCaptor.capture()))
      .thenReturn(null);

  // Simplest schema: a table with int64 key
  preparePkMetadata(tx, Arrays.asList(pkMetadata("tEsT", "key", "ASC")));
  prepareColumnMetadata(tx, Arrays.asList(columnMetadata("tEsT", "key", "INT64", CELLS_PER_KEY)));
}
 

private static void readOnlyTransaction(PrintWriter pw) {
  // ReadOnlyTransaction must be closed by calling close() on it to release resources held by it.
  // We use a try-with-resource block to automatically do so.
  try (ReadOnlyTransaction transaction =
      SpannerClient.getDatabaseClient().readOnlyTransaction()) {
    ResultSet queryResultSet =
        transaction.executeQuery(
            Statement.of("SELECT SingerId, AlbumId, AlbumTitle FROM Albums"));
    while (queryResultSet.next()) {
      pw.printf(
          "%d %d %s\n",
          queryResultSet.getLong(0), queryResultSet.getLong(1), queryResultSet.getString(2));
    }
    ResultSet readResultSet =
        transaction.read(
            "Albums", KeySet.all(), Arrays.asList("SingerId", "AlbumId", "AlbumTitle"));
    while (readResultSet.next()) {
      pw.printf(
          "%d %d %s\n",
          readResultSet.getLong(0), readResultSet.getLong(1), readResultSet.getString(2));
    }
  }
}
 

private static void readOnlyTransaction(PrintWriter pw) {
  // ReadOnlyTransaction must be closed by calling close() on it to release resources held by it.
  // We use a try-with-resource block to automatically do so.
  try (ReadOnlyTransaction transaction =
      SpannerClient.getDatabaseClient().readOnlyTransaction()) {
    ResultSet queryResultSet =
        transaction.executeQuery(
            Statement.of("SELECT SingerId, AlbumId, AlbumTitle FROM Albums"));
    while (queryResultSet.next()) {
      pw.printf(
          "%d %d %s\n",
          queryResultSet.getLong(0), queryResultSet.getLong(1), queryResultSet.getString(2));
    }
    ResultSet readResultSet =
        transaction.read(
            "Albums", KeySet.all(), Arrays.asList("SingerId", "AlbumId", "AlbumTitle"));
    while (readResultSet.next()) {
      pw.printf(
          "%d %d %s\n",
          readResultSet.getLong(0), readResultSet.getLong(1), readResultSet.getString(2));
    }
  }
}
 

private void preparePkMetadata(ReadOnlyTransaction tx, List<Struct> rows) {
  Type type =
      Type.struct(
          Type.StructField.of("table_name", Type.string()),
          Type.StructField.of("column_name", Type.string()),
          Type.StructField.of("column_ordering", Type.string()));
  when(tx.executeQuery(
          argThat(
              new ArgumentMatcher<Statement>() {

                @Override
                public boolean matches(Statement argument) {
                  if (!(argument instanceof Statement)) {
                    return false;
                  }
                  Statement st = (Statement) argument;
                  return st.getSql().contains("information_schema.index_columns");
                }
              })))
      .thenReturn(ResultSets.forRows(type, rows));
}
 

private void prepareColumnMetadata(ReadOnlyTransaction tx, List<Struct> rows) {
  Type type =
      Type.struct(
          Type.StructField.of("table_name", Type.string()),
          Type.StructField.of("column_name", Type.string()),
          Type.StructField.of("spanner_type", Type.string()),
          Type.StructField.of("cells_mutated", Type.int64()));
  when(tx.executeQuery(
          argThat(
              new ArgumentMatcher<Statement>() {

                @Override
                public boolean matches(Statement argument) {
                  if (!(argument instanceof Statement)) {
                    return false;
                  }
                  Statement st = (Statement) argument;
                  return st.getSql().contains("information_schema.columns");
                }
              })))
      .thenReturn(ResultSets.forRows(type, rows));
}
 

static void readOnlyTransaction(DatabaseClient dbClient) {
  // ReadOnlyTransaction must be closed by calling close() on it to release resources held by it.
  // We use a try-with-resource block to automatically do so.
  try (ReadOnlyTransaction transaction = dbClient.readOnlyTransaction()) {
    ResultSet queryResultSet =
        transaction.executeQuery(
            Statement.of("SELECT SingerId, AlbumId, AlbumTitle FROM Albums"));
    while (queryResultSet.next()) {
      System.out.printf(
          "%d %d %s\n",
          queryResultSet.getLong(0), queryResultSet.getLong(1), queryResultSet.getString(2));
    }
    try (ResultSet readResultSet =
        transaction.read(
            "Albums", KeySet.all(), Arrays.asList("SingerId", "AlbumId", "AlbumTitle"))) {
      while (readResultSet.next()) {
        System.out.printf(
            "%d %d %s\n",
            readResultSet.getLong(0), readResultSet.getLong(1), readResultSet.getString(2));
      }
    }
  }
}
 

private ResultSet readTableInfo(ReadOnlyTransaction tx) {
  // retrieve schema information for all tables, as well as aggregating the
  // number of indexes that cover each column. this will be used to estimate
  // the number of cells (table column plus indexes) mutated in an upsert operation
  // in order to stay below the 20k threshold
  return tx.executeQuery(
      Statement.of(
          "SELECT"
              + "    c.table_name"
              + "  , c.column_name"
              + "  , c.spanner_type"
              + "  , (1 + COALESCE(t.indices, 0)) AS cells_mutated"
              + "  FROM ("
              + "    SELECT c.table_name, c.column_name, c.spanner_type, c.ordinal_position"
              + "     FROM information_schema.columns as c"
              + "     WHERE c.table_catalog = '' AND c.table_schema = '') AS c"
              + "  LEFT OUTER JOIN ("
              + "    SELECT t.table_name, t.column_name, COUNT(*) AS indices"
              + "      FROM information_schema.index_columns AS t "
              + "      WHERE t.index_name != 'PRIMARY_KEY' AND t.table_catalog = ''"
              + "      AND t.table_schema = ''"
              + "      GROUP BY t.table_name, t.column_name) AS t"
              + "  USING (table_name, column_name)"
              + "  ORDER BY c.table_name, c.ordinal_position"));
}
 

@Test
public void readOnlyTransactionDmlTest() {

	this.expectedException.expectMessage("A read-only transaction template cannot execute DML.");

	ReadOnlyTransaction readOnlyTransaction = mock(ReadOnlyTransaction.class);
	when(this.databaseClient.readOnlyTransaction(
			// exact staleness is expected.
			eq(TimestampBound.ofReadTimestamp(Timestamp.ofTimeMicroseconds(333)))))
					.thenReturn(readOnlyTransaction);

	this.spannerTemplate
			.performReadOnlyTransaction((spannerOperations) -> {
				spannerOperations.executeDmlStatement(Statement.of("fail"));
				return null;
			}, new SpannerReadOptions()
					.setTimestamp(Timestamp.ofTimeMicroseconds(333)));
}
 

@Test
public void readOnlyTransactionTest() {

	ReadOnlyTransaction readOnlyTransaction = mock(ReadOnlyTransaction.class);
	when(this.databaseClient.readOnlyTransaction(
			eq(TimestampBound.ofMinReadTimestamp(Timestamp.ofTimeMicroseconds(333)))))
					.thenReturn(readOnlyTransaction);

	String finalResult = this.spannerTemplate
			.performReadOnlyTransaction((spannerOperations) -> {
				List<TestEntity> items = spannerOperations.readAll(TestEntity.class);
				TestEntity item = spannerOperations.read(TestEntity.class,
						Key.of("key"));
				return "all done";
			}, new SpannerReadOptions()
					.setTimestampBound(TimestampBound.ofMinReadTimestamp(Timestamp.ofTimeMicroseconds(333L))));

	assertThat(finalResult).isEqualTo("all done");
	verify(readOnlyTransaction, times(2)).read(eq("custom_test_table"), any(), any());
}
 

@Override
public <T> T performReadOnlyTransaction(Function<SpannerTemplate, T> operations,
		SpannerReadOptions readOptions) {
	return doWithOrWithoutTransactionContext((x) -> {
		throw new IllegalStateException("There is already declarative transaction open. " +
				"Spanner does not support nested transactions");
	}, () -> {

		SpannerReadOptions options = (readOptions != null) ? readOptions : new SpannerReadOptions();
		try (ReadOnlyTransaction readOnlyTransaction = (options.getTimestampBound() != null)
				? this.databaseClientProvider.get().readOnlyTransaction(options.getTimestampBound())
				: this.databaseClientProvider.get().readOnlyTransaction()) {
			return operations.apply(new ReadOnlyTransactionSpannerTemplate(
					SpannerTemplate.this.databaseClientProvider,
					SpannerTemplate.this.mappingContext,
					SpannerTemplate.this.spannerEntityProcessor,
					SpannerTemplate.this.mutationFactory,
					SpannerTemplate.this.spannerSchemaUtils, readOnlyTransaction));
		}
	});
}
 

/** Example of read only transaction. */
// [TARGET readOnlyTransaction()]
// [VARIABLE my_singer_id]
// [VARIABLE my_album_id]
public String readOnlyTransaction(long singerId, long albumId) {
  // [START readOnlyTransaction]
  String singerColumn = "FirstName";
  String albumColumn = "AlbumTitle";
  String albumTitle = null;
  // ReadOnlyTransaction should be closed to prevent resource leak.
  try (ReadOnlyTransaction txn = dbClient.readOnlyTransaction()) {
    Struct singerRow =
        txn.readRow("Singers", Key.of(singerId), Collections.singleton(singerColumn));
    Struct albumRow =
        txn.readRow("Albums", Key.of(singerId, albumId), Collections.singleton(albumColumn));
    singerRow.getString(singerColumn);
    albumTitle = albumRow.getString(albumColumn);
  }
  // [END readOnlyTransaction]
  return albumTitle;
}
 

/** Example of read only transaction with timestamp bound. */
// [TARGET readOnlyTransaction(TimestampBound)]
// [VARIABLE my_singer_id]
// [VARIABLE my_album_id]
public String readOnlyTransactionTimestamp(long singerId, long albumId) {
  // [START readOnlyTransactionTimestamp]
  String singerColumn = "FirstName";
  String albumColumn = "AlbumTitle";
  String albumTitle = null;
  // ReadOnlyTransaction should be closed to prevent resource leak.
  try (ReadOnlyTransaction txn =
      dbClient.readOnlyTransaction(TimestampBound.ofExactStaleness(10, TimeUnit.SECONDS))) {
    Struct singerRow =
        txn.readRow("Singers", Key.of(singerId), Collections.singleton(singerColumn));
    Struct albumRow =
        txn.readRow("Albums", Key.of(singerId, albumId), Collections.singleton(albumColumn));
    singerRow.getString(singerColumn);
    albumTitle = albumRow.getString(albumColumn);
  }
  // [END readOnlyTransactionTimestamp]
  return albumTitle;
}
 

@ProcessElement
@SuppressWarnings("unused")
public void processElement(ProcessContext processContext) {
  // Save schema to GCS so it can be saved along with the exported file.
  LOG.info("Creating database client for schema read");

  LinkedHashMap<String, String> columns;
  try {
    DatabaseClient databaseClient = getDatabaseClient(spannerConfig());

    try (ReadOnlyTransaction context = databaseClient.readOnlyTransaction()) {
      LOG.info("Reading schema information");
      columns = getAllColumns(context, table().get());
      String columnJson = SpannerConverters.GSON.toJson(columns);
      LOG.info("Saving schema information");
      saveSchema(columnJson, textWritePrefix().get() + SCHEMA_SUFFIX);
    }
  } finally {
    closeSpannerAccessor();
  }

  processContext.output(
      ReadOperation.create()
          .withColumns(new ArrayList<>(columns.keySet()))
          .withTable(table().get()));
}
 

private ResultSet readTableInfo(ReadOnlyTransaction tx) {
  // retrieve schema information for all tables, as well as aggregating the
  // number of indexes that cover each column. this will be used to estimate
  // the number of cells (table column plus indexes) mutated in an upsert operation
  // in order to stay below the 20k threshold
  return tx.executeQuery(
      Statement.of(
          "SELECT"
              + "    c.table_name"
              + "  , c.column_name"
              + "  , c.spanner_type"
              + "  , (1 + COALESCE(t.indices, 0)) AS cells_mutated"
              + "  FROM ("
              + "    SELECT c.table_name, c.column_name, c.spanner_type, c.ordinal_position"
              + "     FROM information_schema.columns as c"
              + "     WHERE c.table_catalog = '' AND c.table_schema = '') AS c"
              + "  LEFT OUTER JOIN ("
              + "    SELECT t.table_name, t.column_name, COUNT(*) AS indices"
              + "      FROM information_schema.index_columns AS t "
              + "      WHERE t.index_name != 'PRIMARY_KEY' AND t.table_catalog = ''"
              + "      AND t.table_schema = ''"
              + "      GROUP BY t.table_name, t.column_name) AS t"
              + "  USING (table_name, column_name)"
              + "  ORDER BY c.table_name, c.ordinal_position"));
}
 

public void read(ReadQuery request, StreamObserver<EmptyResponse> responseObserver) {
  System.out.println("read has been called");

  try (ReadOnlyTransaction transaction = dbClient.readOnlyTransaction()) {
    ResultSet resultSet = transaction.executeQuery(Statement.of(request.getQuery()));
    while (resultSet.next()) {
      // Do nothing with the data.
    }
  }

  EmptyResponse reply = EmptyResponse.newBuilder().build();
  responseObserver.onNext(reply);
  responseObserver.onCompleted();
}
 

/** Example of single use read only transaction. */
// [TARGET singleUseReadOnlyTransaction()]
// [VARIABLE my_singer_id]
public Timestamp singleUseReadOnlyTransaction(long singerId) {
  // [START singleUseReadOnlyTransaction]
  String column = "FirstName";
  ReadOnlyTransaction txn = dbClient.singleUseReadOnlyTransaction();
  Struct row = txn.readRow("Singers", Key.of(singerId), Collections.singleton(column));
  row.getString(column);
  Timestamp timestamp = txn.getReadTimestamp();
  // [END singleUseReadOnlyTransaction]
  return timestamp;
}
 

/** Example of single use read only transaction with timestamp bound. */
// [TARGET singleUseReadOnlyTransaction(TimestampBound)]
// [VARIABLE my_singer_id]
public Timestamp singleUseReadOnlyTransactionTimestamp(long singerId) {
  // [START singleUseReadOnlyTransactionTimestamp]
  String column = "FirstName";
  ReadOnlyTransaction txn =
      dbClient.singleUseReadOnlyTransaction(TimestampBound.ofMaxStaleness(10, TimeUnit.SECONDS));
  Struct row = txn.readRow("Singers", Key.of(singerId), Collections.singleton(column));
  row.getString(column);
  Timestamp timestamp = txn.getReadTimestamp();
  // [END singleUseReadOnlyTransactionTimestamp]
  return timestamp;
}
 

@Test
public void noBatching() throws Exception {

  // This test uses a different mock/fake because it explicitly does not want to populate the
  // Spanner schema.
  FakeServiceFactory fakeServiceFactory = new FakeServiceFactory();
  ReadOnlyTransaction tx = mock(ReadOnlyTransaction.class);
  when(fakeServiceFactory.mockDatabaseClient().readOnlyTransaction()).thenReturn(tx);

  // Capture batches sent to writeAtLeastOnce.
  when(fakeServiceFactory.mockDatabaseClient().writeAtLeastOnce(mutationBatchesCaptor.capture()))
      .thenReturn(null);

  PCollection<MutationGroup> mutations = pipeline.apply(Create.of(g(m(1L)), g(m(2L))));
  mutations.apply(
      SpannerIO.write()
          .withProjectId("test-project")
          .withInstanceId("test-instance")
          .withDatabaseId("test-database")
          .withServiceFactory(fakeServiceFactory)
          .withBatchSizeBytes(1)
          .grouped());
  pipeline.run();

  verify(fakeServiceFactory.mockDatabaseClient(), times(1))
      .writeAtLeastOnce(mutationsInNoOrder(batch(m(1L))));
  verify(fakeServiceFactory.mockDatabaseClient(), times(1))
      .writeAtLeastOnce(mutationsInNoOrder(batch(m(2L))));
  // If no batching then the DB schema is never read.
  verify(tx, never()).executeQuery(any());
}
 

@Test
public void simple() throws Exception {
  // Simplest schema: a table with int64 key
  ReadOnlyTransaction tx = mock(ReadOnlyTransaction.class);
  when(serviceFactory.mockDatabaseClient().readOnlyTransaction()).thenReturn(tx);

  preparePkMetadata(tx, Arrays.asList(pkMetadata("test", "key", "ASC")));
  prepareColumnMetadata(tx, Arrays.asList(columnMetadata("test", "key", "INT64")));

  SpannerConfig config =
      SpannerConfig.create()
          .withProjectId("test-project")
          .withInstanceId("test-instance")
          .withDatabaseId("test-database")
          .withServiceFactory(serviceFactory);

  DoFnTester<Void, SpannerSchema> tester = DoFnTester.of(new ReadSpannerSchema(config));
  List<SpannerSchema> schemas = tester.processBundle(Arrays.asList((Void) null));

  assertEquals(1, schemas.size());

  SpannerSchema schema = schemas.get(0);

  assertEquals(1, schema.getTables().size());

  SpannerSchema.Column column = SpannerSchema.Column.create("key", Type.int64());
  SpannerSchema.KeyPart keyPart = SpannerSchema.KeyPart.create("key", false);

  assertThat(schema.getColumns("test"), contains(column));
  assertThat(schema.getKeyParts("test"), contains(keyPart));
}
 

private ResultSet readPrimaryKeyInfo(ReadOnlyTransaction tx) {
  return tx.executeQuery(
      Statement.of(
          "SELECT t.table_name, t.column_name, t.column_ordering"
              + " FROM information_schema.index_columns AS t "
              + " WHERE t.index_name = 'PRIMARY_KEY' AND t.table_catalog = ''"
              + " AND t.table_schema = ''"
              + " ORDER BY t.table_name, t.ordinal_position"));
}
 

ReadOnlyTransactionSpannerTemplate(Supplier<DatabaseClient> databaseClient,
		SpannerMappingContext mappingContext,
		SpannerEntityProcessor spannerEntityProcessor,
		SpannerMutationFactory spannerMutationFactory,
		SpannerSchemaUtils spannerSchemaUtils,
		ReadOnlyTransaction readOnlyTransaction) {
	super(databaseClient, mappingContext, spannerEntityProcessor,
			spannerMutationFactory, spannerSchemaUtils);
	this.readOnlyTransaction = readOnlyTransaction;
}
 

private Ddl readDdl(String db) {
  DatabaseClient dbClient = spannerServer.getDbClient(db);
  Ddl ddl;
  try (ReadOnlyTransaction ctx = dbClient.readOnlyTransaction()) {
    ddl = new InformationSchemaScanner(ctx).scan();
  }
  return ddl;
}
 

private Ddl readDdl(String db) {
  SpannerOptions spannerOptions = SpannerOptions.newBuilder().build();
  Spanner client = spannerOptions.getService();
  Ddl ddl;
  try (ReadOnlyTransaction ctx = spannerServer.getDbClient(db).readOnlyTransaction()) {
    ddl = new InformationSchemaScanner(ctx).scan();
  }
  return ddl;
}
 

/** Unit test for export transform. */
@Test
@Category(NeedsRunner.class)
public void testSchemaSave() throws IOException {

  ValueProvider<String> table = ValueProvider.StaticValueProvider.of(TABLE);
  SpannerConfig spannerConfig = SpannerConfig.create();
  DatabaseClient databaseClient = mock(DatabaseClient.class, withSettings().serializable());
  ReadOnlyTransaction readOnlyTransaction =
      mock(ReadOnlyTransaction.class, withSettings().serializable());
  ResultSet resultSet = mock(ResultSet.class, withSettings().serializable());
  Struct struct = mock(Struct.class, withSettings().serializable());

  when(databaseClient.readOnlyTransaction()).thenReturn(readOnlyTransaction);
  when(readOnlyTransaction.executeQuery(any(Statement.class))).thenReturn(resultSet);
  when(resultSet.next()).thenReturn(true).thenReturn(false);
  when(resultSet.getCurrentRowAsStruct()).thenReturn(struct);
  when(struct.getString(0)).thenReturn(COLUMN_NAME);
  when(struct.getString(1)).thenReturn("INT64");

  String schemaPath = "/tmp/" + UUID.randomUUID().toString();
  ValueProvider<String> textWritePrefix = ValueProvider.StaticValueProvider.of(schemaPath);
  SpannerConverters.ExportTransform exportTransform =
      SpannerConverters.ExportTransformFactory.create(table, spannerConfig, textWritePrefix);
  exportTransform.setDatabaseClient(databaseClient);

  PCollection<ReadOperation> results = pipeline.apply("Create", exportTransform);
  ReadOperation readOperation =
      ReadOperation.create().withColumns(ImmutableList.of(COLUMN_NAME)).withTable(TABLE);
  PAssert.that(results).containsInAnyOrder(readOperation);
  pipeline.run();
  ReadableByteChannel channel =
      FileSystems.open(
          FileSystems.matchNewResource(
              schemaPath + SpannerConverters.ExportTransform.ExportFn.SCHEMA_SUFFIX,
              false));
  java.util.Scanner scanner = new java.util.Scanner(channel).useDelimiter("\\A");
  assertEquals("{\"id\":\"INT64\"}", scanner.next());
}
 

private ResultSet readPrimaryKeyInfo(ReadOnlyTransaction tx) {
  return tx.executeQuery(
      Statement.of(
          "SELECT t.table_name, t.column_name, t.column_ordering"
              + " FROM information_schema.index_columns AS t "
              + " WHERE t.index_name = 'PRIMARY_KEY' AND t.table_catalog = ''"
              + " AND t.table_schema = ''"
              + " ORDER BY t.table_name, t.ordinal_position"));
}
 

@Before
@SuppressWarnings("unchecked")
public void setUp() throws Exception {
  serviceFactory = new FakeServiceFactory();
  mockTx = mock(ReadOnlyTransaction.class);
}