com.google.cloud.spanner.ReadContext Java Examples

/** Function to get all column names from the table. */
private LinkedHashMap<String, String> getAllColumns(ReadContext context, String tableName) {
  LinkedHashMap<String, String> columns = Maps.newLinkedHashMap();
  ResultSet resultSet =
      context.executeQuery(
          Statement.newBuilder(
                  "SELECT COLUMN_NAME, SPANNER_TYPE FROM INFORMATION_SCHEMA.COLUMNS "
                      + "WHERE TABLE_NAME=@table_name ORDER BY ORDINAL_POSITION")
              .bind("table_name")
              .to(tableName)
              .build());
  LOG.info("Got schema information. Reading columns.");
  while (resultSet.next()) {
    Struct currentRow = resultSet.getCurrentRowAsStruct();
    columns.put(currentRow.getString(0), currentRow.getString(1));
  }
  return columns;
}
 

@Override
public ResultSet executeQuery() throws SQLException {
  CustomDriverStatement custom = getCustomDriverStatement(sqlTokens);
  if (custom != null && custom.isQuery()) {
    return custom.executeQuery(sqlTokens);
  }
  Statement statement;
  try {
    statement = CCJSqlParserUtil.parse(sanitizeSQL(sql));
  } catch (JSQLParserException | TokenMgrException e) {
    throw new CloudSpannerSQLException(PARSE_ERROR + sql + ": " + e.getLocalizedMessage(),
        Code.INVALID_ARGUMENT, e);
  }
  if (statement instanceof Select) {
    determineForceSingleUseReadContext((Select) statement);
    com.google.cloud.spanner.Statement.Builder builder = createSelectBuilder(statement, sql);
    try (ReadContext context = getReadContext()) {
      com.google.cloud.spanner.ResultSet rs = context.executeQuery(builder.build());
      return new CloudSpannerResultSet(this, rs, sql);
    }
  }
  throw new CloudSpannerSQLException(
      "SQL statement not suitable for executeQuery. Expected SELECT-statement.",
      Code.INVALID_ARGUMENT);
}
 

private ResultSet executeRead(String tableName, KeySet keys, Iterable<String> columns,
		SpannerReadOptions options) {

	long startTime = LOGGER.isDebugEnabled() ? System.currentTimeMillis() : 0;

	ReadContext readContext = (options != null && options.getTimestampBound() != null)
			? getReadContext(options.getTimestampBound())
			: getReadContext();

	final ResultSet resultSet = options != null && options.getIndex() != null
			? readContext.readUsingIndex(tableName, options.getIndex(), keys, columns, options.getOptions())
			: readContext.read(tableName, keys, columns, options == null ? ArrayUtils.toArray() : options.getOptions());

	if (LOGGER.isDebugEnabled()) {
		StringBuilder logs = logColumns(tableName, keys, columns);
		logReadOptions(options, logs);
		LOGGER.debug(logs.toString());

		LOGGER.debug("Read elapsed milliseconds: " + (System.currentTimeMillis() - startTime));
	}

	return resultSet;
}
 

ResultSetStats analyzeQuery() {
  // [START read_context_analyze_query]
  ReadContext rc = dbClient.singleUse();
  ResultSet resultSet =
      rc.analyzeQuery(
          Statement.of("SELECT SingerId, AlbumId, MarketingBudget FROM Albums"),
          ReadContext.QueryAnalyzeMode.PROFILE);
  while (resultSet.next()) {
    // Discard the results. We're only processing because getStats() below requires it.
    resultSet.getCurrentRowAsStruct();
  }
  ResultSetStats stats = resultSet.getStats();
  // [END read_context_analyze_query]
  return stats;
}
 

ResultSet executeQuery() {
  // [START read_context_execute_query]
  // Rows without an explicit value for MarketingBudget will have a MarketingBudget equal to
  // null.
  ReadContext readContext = dbClient.singleUse();
  ResultSet resultSet =
      readContext.executeQuery(
          Statement.of("SELECT SingerId, AlbumId, MarketingBudget, LastUpdateTime FROM Albums"));
  // [END read_context_execute_query]
  return resultSet;
}
 

@Override
public ResultSet executeQuery(String sql) throws SQLException {
  String[] sqlTokens = getTokens(sql);
  CustomDriverStatement custom = getCustomDriverStatement(sqlTokens);
  if (custom != null && custom.isQuery()) {
    return custom.executeQuery(sqlTokens);
  }
  try (ReadContext context = getReadContext()) {
    com.google.cloud.spanner.ResultSet rs =
        context.executeQuery(com.google.cloud.spanner.Statement.of(sql));
    return new CloudSpannerResultSet(this, rs, sql);
  }
}
 

Struct readRowUsingIndex() {
  // [START read_context_read_index]
  ReadContext readContext = dbClient.singleUse();
  Struct row =
      readContext.readRowUsingIndex(
          "Albums",
          "AlbumsByAlbumId",
          Key.of(1, "Green"),
          Arrays.asList("AlbumId", "AlbumTitle"));
  // [END read_context_read_index]

  return row;
}
 

Struct readRow() {
  // [START read_context_read_row]
  ReadContext readContext = dbClient.singleUse();
  Struct row = readContext.readRow("Albums", Key.of(2, 1), Arrays.asList("MarketingBudget"));
  // [END read_context_read_row]

  return row;
}
 

ResultSet readUsingIndex() {
  // [START read_context_read_index]
  ReadContext readContext = dbClient.singleUse();
  ResultSet resultSet =
      readContext.readUsingIndex(
          "Albums", "AlbumsByAlbumTitle", KeySet.all(), Arrays.asList("AlbumId", "AlbumTitle"));
  // [END read_context_read_index]

  return resultSet;
}
 

ResultSet read() {
  // [START read_context_read]
  ReadContext readContext = dbClient.singleUse();
  ResultSet resultSet =
      readContext.read(
          "Albums",
          // KeySet.all() can be used to read all rows in a table. KeySet exposes other
          // methods to read only a subset of the table.
          KeySet.all(),
          Arrays.asList("SingerId", "AlbumId", "AlbumTitle"));
  // [END read_context_read]

  return resultSet;
}
 

private void setupFailedMockQuery() {
  ReadContext readContext = mock(ReadContext.class);
  when(readContext.executeQuery(Statement.of("SELECT SingerId, AlbumId, AlbumTitle FROM Albums")))
      .thenThrow(
          SpannerExceptionFactory.newSpannerException(
              ErrorCode.NOT_FOUND, "Table `Albums` not found"));
  when(client.singleUse()).thenReturn(readContext);
}
 

private void setupSuccessfulMockQuery() {
  ReadContext readContext = mock(ReadContext.class);
  ResultSet resultSet = mock(ResultSet.class);
  when(resultSet.next()).thenReturn(true, true, false);
  when(resultSet.getLong("SingerId")).thenReturn(1L, 2L, 0L);
  when(resultSet.getLong("AlbumId")).thenReturn(1L, 1L, 0L);
  when(resultSet.getString("AlbumTitle")).thenReturn("Album 1", "Album 2", null);
  when(readContext.executeQuery(Statement.of("SELECT SingerId, AlbumId, AlbumTitle FROM Albums")))
      .thenReturn(resultSet);
  when(client.singleUse()).thenReturn(readContext);
}
 

@Before
public void setUp() {
	this.databaseClient = mock(DatabaseClient.class);
	this.mappingContext = new SpannerMappingContext();
	this.objectMapper = mock(SpannerEntityProcessor.class);
	this.mutationFactory = mock(SpannerMutationFactory.class);
	this.schemaUtils = new SpannerSchemaUtils(this.mappingContext, this.objectMapper,
			true);
	this.readContext = mock(ReadContext.class);
	when(this.databaseClient.singleUse()).thenReturn(this.readContext);
	when(this.objectMapper.getWriteConverter()).thenReturn(new SpannerWriteConverter());
	this.spannerTemplate = new SpannerTemplate(() -> this.databaseClient,
			this.mappingContext, this.objectMapper, this.mutationFactory,
			this.schemaUtils);
}
 

@Override
protected ReadContext getReadContext() {
	return this.transactionContext;
}
 

@Test
public void getTableRelationshipsTest() {
	ReadContext readContext = mock(ReadContext.class);

	Struct s1 = Struct.newBuilder().set("table_name").to(Value.string("grandpa"))
			.set("parent_table_name").to(Value.string(null)).build();
	Struct s2 = Struct.newBuilder().set("table_name").to(Value.string("parent_a"))
			.set("parent_table_name").to(Value.string("grandpa")).build();
	Struct s3 = Struct.newBuilder().set("table_name").to(Value.string("parent_b"))
			.set("parent_table_name").to(Value.string("grandpa")).build();
	Struct s4 = Struct.newBuilder().set("table_name").to(Value.string("child"))
			.set("parent_table_name").to(Value.string("parent_a")).build();

	MockResults mockResults = new MockResults();
	mockResults.structs = Arrays.asList(s1, s2, s3, s4);
	ResultSet results = mock(ResultSet.class);
	when(results.next()).thenAnswer((invocation) -> mockResults.next());
	when(results.getCurrentRowAsStruct())
			.thenAnswer((invocation) -> mockResults.getCurrent());
	when(this.databaseClient.singleUse()).thenReturn(readContext);
	when(readContext.executeQuery(any())).thenReturn(results);

	Map<String, Set<String>> relationships = this.spannerDatabaseAdminTemplate
			.getParentChildTablesMap();

	assertThat(relationships).hasSize(2);
	assertThat(relationships.get("grandpa")).containsExactlyInAnyOrder("parent_a", "parent_b");
	assertThat(relationships.get("parent_a")).containsExactlyInAnyOrder("child");

	assertThat(this.spannerDatabaseAdminTemplate.isInterleaved("grandpa", "child"))
			.as("verify grand-child relationship").isTrue();
	assertThat(this.spannerDatabaseAdminTemplate.isInterleaved("grandpa", "parent_a"))
			.as("verify parent-child relationship").isTrue();
	assertThat(this.spannerDatabaseAdminTemplate.isInterleaved("parent_a", "child"))
			.as("verify parent-child relationship").isTrue();
	assertThat(this.spannerDatabaseAdminTemplate.isInterleaved("grandpa", "parent_b"))
			.as("verify parent-child relationship").isTrue();
	assertThat(this.spannerDatabaseAdminTemplate.isInterleaved("parent_a", "parent_b"))
			.as("verify not parent-child relationship").isFalse();
	assertThat(this.spannerDatabaseAdminTemplate.isInterleaved("parent_b", "child"))
			.as("verify not parent-child relationship").isFalse();
}
 

@Override
protected ReadContext getReadContext(TimestampBound timestampBound) {
	throw new SpannerDataException(
			"Getting stale snapshot read contexts is not supported"
					+ " in read-write transaction templates.");
}
 

@Override
protected ReadContext getReadContext(TimestampBound timestampBound) {
	throw new SpannerDataException(
			"Getting stale snapshot read contexts is not supported"
					+ " in read-only transaction templates.");
}
 

@Override
protected ReadContext getReadContext() {
	return this.readOnlyTransaction;
}
 

protected ReadContext getReadContext(TimestampBound timestampBound) {
	return doWithOrWithoutTransactionContext((x) -> x,
			() -> this.databaseClientProvider.get().singleUse(timestampBound));
}
 

protected ReadContext getReadContext() {
	return doWithOrWithoutTransactionContext((x) -> x, this.databaseClientProvider.get()::singleUse);
}
 

@Override
public boolean execute() throws SQLException {
  CustomDriverStatement custom = getCustomDriverStatement(sqlTokens);
  if (custom != null)
    return custom.execute(sqlTokens);
  Statement statement = null;
  boolean ddl = isDDLStatement();
  if (!ddl) {
    try {
      statement = CCJSqlParserUtil.parse(sanitizeSQL(sql));
    } catch (JSQLParserException | TokenMgrException e) {
      throw new CloudSpannerSQLException(PARSE_ERROR + sql + ": " + e.getLocalizedMessage(),
          Code.INVALID_ARGUMENT, e);
    }
  }
  if (!ddl && statement instanceof Select) {
    determineForceSingleUseReadContext((Select) statement);
    com.google.cloud.spanner.Statement.Builder builder = createSelectBuilder(statement, sql);
    if (!isForceSingleUseReadContext() && getConnection().isBatchReadOnly()) {
      List<Partition> partitions = partitionQuery(builder.build());
      currentResultSets = new ArrayList<>(partitions.size());
      for (Partition p : partitions) {
        currentResultSets
            .add(new CloudSpannerPartitionResultSet(this, getBatchReadOnlyTransaction(), p, sql));
      }
    } else {
      try (ReadContext context = getReadContext()) {
        com.google.cloud.spanner.ResultSet rs = context.executeQuery(builder.build());
        currentResultSets = Arrays.asList(new CloudSpannerResultSet(this, rs, sql));
        currentResultSetIndex = 0;
        lastUpdateCount = -1;
      }
    }
    return true;
  } else {
    lastUpdateCount = executeUpdate();
    currentResultSets = null;
    currentResultSetIndex = 0;
    return false;
  }
}
 

@Override
public boolean execute(String sql) throws SQLException {
  String[] sqlTokens = getTokens(sql);
  CustomDriverStatement custom = getCustomDriverStatement(sqlTokens);
  if (custom != null)
    return custom.execute(sqlTokens);
  Statement statement = null;
  boolean ddl = isDDLStatement(sqlTokens);
  if (!ddl) {
    try {
      statement = CCJSqlParserUtil.parse(sanitizeSQL(sql));
    } catch (JSQLParserException | TokenMgrException e) {
      throw new CloudSpannerSQLException(
          "Error while parsing sql statement " + sql + ": " + e.getLocalizedMessage(),
          Code.INVALID_ARGUMENT, e);
    }
  }
  if (!ddl && statement instanceof Select) {
    determineForceSingleUseReadContext((Select) statement);
    if (!isForceSingleUseReadContext() && getConnection().isBatchReadOnly()) {
      List<Partition> partitions = partitionQuery(com.google.cloud.spanner.Statement.of(sql));
      currentResultSets = new ArrayList<>(partitions.size());
      for (Partition p : partitions) {
        currentResultSets
            .add(new CloudSpannerPartitionResultSet(this, getBatchReadOnlyTransaction(), p, sql));
      }
    } else {
      try (ReadContext context = getReadContext()) {
        com.google.cloud.spanner.ResultSet rs =
            context.executeQuery(com.google.cloud.spanner.Statement.of(sql));
        currentResultSets = Arrays.asList(new CloudSpannerResultSet(this, rs, sql));
        currentResultSetIndex = 0;
        lastUpdateCount = -1;
      }
    }
    return true;
  } else {
    lastUpdateCount = executeUpdate(sql);
    currentResultSetIndex = 0;
    currentResultSets = null;
    return false;
  }
}
 

protected ReadContext getReadContext() throws SQLException {
  if (connection.getAutoCommit() || forceSingleUseReadContext) {
    return dbClient.singleUse();
  }
  return connection.getTransaction();
}
 

public InformationSchemaScanner(ReadContext context) {
  this.context = context;
}