Skip to content

Commit b6f2949

Browse files
nareshznareshz
committed
Add tests for new queryMode's
1 parent 465ab61 commit b6f2949

File tree

8 files changed

+245
-7
lines changed

8 files changed

+245
-7
lines changed

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/Connection.java‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1327,9 +1327,13 @@ PartitionedQueryResultSet runPartitionedQuery(
13271327
* Analyzes a DML statement and returns query plan and/or execution statistics information.
13281328
*
13291329
* <p>{@link com.google.cloud.spanner.ReadContext.QueryAnalyzeMode#PLAN} only returns the plan for
1330-
* the statement. {@link com.google.cloud.spanner.ReadContext.QueryAnalyzeMode#PROFILE} executes
1331-
* the DML statement, returns the modified row count and execution statistics, and the effects of
1332-
* the DML statement will be visible to subsequent operations in the transaction.
1330+
* the statement. {@link com.google.cloud.spanner.ReadContext.QueryAnalyzeMode#WITH_STATS} returns
1331+
* the overall (but not operator-level) execution statistics. {@link
1332+
* com.google.cloud.spanner.ReadContext.QueryAnalyzeMode#WITH_PLAN_AND_STATS} returns the query
1333+
* plan and overall (but not operator-level) execution statistics. {@link
1334+
* com.google.cloud.spanner.ReadContext.QueryAnalyzeMode#PROFILE} executes the DML statement,
1335+
* returns the modified row count and execution statistics, and the effects of the DML statement
1336+
* will be visible to subsequent operations in the transaction.
13331337
*
13341338
* @deprecated Use {@link #analyzeUpdateStatement(Statement, QueryAnalyzeMode, UpdateOption...)}
13351339
* instead
@@ -1345,6 +1349,10 @@ default ResultSetStats analyzeUpdate(Statement update, QueryAnalyzeMode analyzeM
13451349
*
13461350
* <p>{@link com.google.cloud.spanner.ReadContext.QueryAnalyzeMode#PLAN} only returns the plan and
13471351
* undeclared parameters for the statement. {@link
1352+
* com.google.cloud.spanner.ReadContext.QueryAnalyzeMode#WITH_STATS} returns the overall (but not
1353+
* operator-level) execution statistics. {@link
1354+
* com.google.cloud.spanner.ReadContext.QueryAnalyzeMode#WITH_PLAN_AND_STATS} returns the query
1355+
* plan and overall (but not operator-level) execution statistics. {@link
13481356
* com.google.cloud.spanner.ReadContext.QueryAnalyzeMode#PROFILE} also executes the DML statement,
13491357
* returns the modified row count and execution statistics, and the effects of the DML statement
13501358
* will be visible to subsequent operations in the transaction.

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3555,9 +3555,7 @@ public void testWithStatsQueryModeWithAnalyzeQuery() {
35553555
try (ReadOnlyTransaction tx = client.readOnlyTransaction()) {
35563556
try (ResultSet rs =
35573557
tx.analyzeQuery(
3558-
Statement.newBuilder(SELECT1.getSql())
3559-
.build(),
3560-
QueryAnalyzeMode.WITH_STATS)) {
3558+
Statement.newBuilder(SELECT1.getSql()).build(), QueryAnalyzeMode.WITH_STATS)) {
35613559
// Just iterate over the results to execute the query.
35623560
consumeResults(rs);
35633561
}

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcResultSetTest.java‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -470,6 +470,63 @@ public void planResult() {
470470
resultSet.close();
471471
}
472472

473+
@Test
474+
public void withStatsResult() {
475+
Map<String, com.google.protobuf.Value> statsMap =
476+
ImmutableMap.of(
477+
"f1", Value.string("").toProto(),
478+
"f2", Value.string("").toProto());
479+
ResultSetStats stats =
480+
ResultSetStats.newBuilder()
481+
.setQueryStats(com.google.protobuf.Struct.newBuilder().putAllFields(statsMap).build())
482+
.build();
483+
ArrayList<Type.StructField> dataType = new ArrayList<>();
484+
dataType.add(Type.StructField.of("data", Type.string()));
485+
consumer.onPartialResultSet(
486+
PartialResultSet.newBuilder()
487+
.setMetadata(makeMetadata(Type.struct(dataType)))
488+
.addValues(Value.string("d1").toProto())
489+
.setChunkedValue(false)
490+
.setStats(stats)
491+
.build());
492+
resultSet = resultSetWithMode(QueryMode.WITH_STATS);
493+
consumer.onCompleted();
494+
assertThat(resultSet.next()).isTrue();
495+
assertThat(resultSet.next()).isFalse();
496+
ResultSetStats receivedStats = resultSet.getStats();
497+
assertThat(stats).isEqualTo(receivedStats);
498+
resultSet.close();
499+
}
500+
501+
@Test
502+
public void withPlanAndStatsResult() {
503+
Map<String, com.google.protobuf.Value> statsMap =
504+
ImmutableMap.of(
505+
"f1", Value.string("").toProto(),
506+
"f2", Value.string("").toProto());
507+
ResultSetStats stats =
508+
ResultSetStats.newBuilder()
509+
.setQueryPlan(QueryPlan.newBuilder().build())
510+
.setQueryStats(com.google.protobuf.Struct.newBuilder().putAllFields(statsMap).build())
511+
.build();
512+
ArrayList<Type.StructField> dataType = new ArrayList<>();
513+
dataType.add(Type.StructField.of("data", Type.string()));
514+
consumer.onPartialResultSet(
515+
PartialResultSet.newBuilder()
516+
.setMetadata(makeMetadata(Type.struct(dataType)))
517+
.addValues(Value.string("d1").toProto())
518+
.setChunkedValue(false)
519+
.setStats(stats)
520+
.build());
521+
resultSet = resultSetWithMode(QueryMode.WITH_PLAN_AND_STATS);
522+
consumer.onCompleted();
523+
assertThat(resultSet.next()).isTrue();
524+
assertThat(resultSet.next()).isFalse();
525+
ResultSetStats receivedStats = resultSet.getStats();
526+
assertThat(stats).isEqualTo(receivedStats);
527+
resultSet.close();
528+
}
529+
473530
@Test
474531
public void statsUnavailable() {
475532
ResultSetStats stats = ResultSetStats.newBuilder().build();

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionImplTest.java‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -272,6 +272,11 @@ public static ConnectionImpl createConnection(final ConnectionOptions options, D
272272
.thenReturn(select1ResultSetWithStats);
273273
when(singleUseReadOnlyTx.analyzeQuery(Statement.of(SELECT), QueryAnalyzeMode.PROFILE))
274274
.thenReturn(select1ResultSetWithStats);
275+
when(singleUseReadOnlyTx.analyzeQuery(Statement.of(SELECT), QueryAnalyzeMode.WITH_STATS))
276+
.thenReturn(select1ResultSetWithStats);
277+
when(singleUseReadOnlyTx.analyzeQuery(
278+
Statement.of(SELECT), QueryAnalyzeMode.WITH_PLAN_AND_STATS))
279+
.thenReturn(select1ResultSetWithStats);
275280
when(singleUseReadOnlyTx.getReadTimestamp())
276281
.then(
277282
invocation -> {
@@ -307,6 +312,11 @@ public static ConnectionImpl createConnection(final ConnectionOptions options, D
307312
.thenReturn(select1ResultSetWithStats);
308313
when(txContext.analyzeQuery(Statement.of(SELECT), QueryAnalyzeMode.PROFILE))
309314
.thenReturn(select1ResultSetWithStats);
315+
when(txContext.analyzeQuery(Statement.of(SELECT), QueryAnalyzeMode.WITH_STATS))
316+
.thenReturn(select1ResultSetWithStats);
317+
when(txContext.analyzeQuery(
318+
Statement.of(SELECT), QueryAnalyzeMode.WITH_PLAN_AND_STATS))
319+
.thenReturn(select1ResultSetWithStats);
310320
when(txContext.executeUpdate(Statement.of(UPDATE))).thenReturn(1L);
311321
return new SimpleTransactionManager(txContext, options.isReturnCommitStats());
312322
});
@@ -328,6 +338,10 @@ public static ConnectionImpl createConnection(final ConnectionOptions options, D
328338
.thenReturn(select1ResultSetWithStats);
329339
when(tx.analyzeQuery(Statement.of(SELECT), QueryAnalyzeMode.PROFILE))
330340
.thenReturn(select1ResultSetWithStats);
341+
when(tx.analyzeQuery(Statement.of(SELECT), QueryAnalyzeMode.WITH_STATS))
342+
.thenReturn(select1ResultSetWithStats);
343+
when(tx.analyzeQuery(Statement.of(SELECT), QueryAnalyzeMode.WITH_PLAN_AND_STATS))
344+
.thenReturn(select1ResultSetWithStats);
331345
when(tx.getReadTimestamp())
332346
.then(
333347
ignored -> {

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadOnlyTransactionTest.java‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -358,6 +358,54 @@ public void testPlanQuery() {
358358
}
359359
}
360360

361+
@Test
362+
public void testWithStatsQuery() {
363+
for (TimestampBound staleness : getTestTimestampBounds()) {
364+
ParsedStatement parsedStatement = mock(ParsedStatement.class);
365+
when(parsedStatement.getType()).thenReturn(StatementType.QUERY);
366+
when(parsedStatement.isQuery()).thenReturn(true);
367+
Statement statement = Statement.of("SELECT * FROM FOO");
368+
when(parsedStatement.getStatement()).thenReturn(statement);
369+
when(parsedStatement.getSqlWithoutComments()).thenReturn(statement.getSql());
370+
371+
ReadOnlyTransaction transaction = createSubject(staleness);
372+
ResultSet rs =
373+
get(
374+
transaction.executeQueryAsync(
375+
CallType.SYNC, parsedStatement, AnalyzeMode.WITH_STATS));
376+
assertThat(rs, is(notNullValue()));
377+
// get all results and then get the stats
378+
while (rs.next()) {
379+
// do nothing
380+
}
381+
assertThat(rs.getStats(), is(notNullValue()));
382+
}
383+
}
384+
385+
@Test
386+
public void testWithPlanAndStatsQuery() {
387+
for (TimestampBound staleness : getTestTimestampBounds()) {
388+
ParsedStatement parsedStatement = mock(ParsedStatement.class);
389+
when(parsedStatement.getType()).thenReturn(StatementType.QUERY);
390+
when(parsedStatement.isQuery()).thenReturn(true);
391+
Statement statement = Statement.of("SELECT * FROM FOO");
392+
when(parsedStatement.getStatement()).thenReturn(statement);
393+
when(parsedStatement.getSqlWithoutComments()).thenReturn(statement.getSql());
394+
395+
ReadOnlyTransaction transaction = createSubject(staleness);
396+
ResultSet rs =
397+
get(
398+
transaction.executeQueryAsync(
399+
CallType.SYNC, parsedStatement, AnalyzeMode.WITH_PLAN_AND_STATS));
400+
assertThat(rs, is(notNullValue()));
401+
// get all results and then get the stats
402+
while (rs.next()) {
403+
// do nothing
404+
}
405+
assertThat(rs.getStats(), is(notNullValue()));
406+
}
407+
}
408+
361409
@Test
362410
public void testProfileQuery() {
363411
for (TimestampBound staleness : getTestTimestampBounds()) {

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadWriteTransactionTest.java‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -268,6 +268,42 @@ public void testProfileQuery() {
268268
assertThat(rs.getStats(), is(notNullValue()));
269269
}
270270

271+
@Test
272+
public void testWithStatsQuery() {
273+
ParsedStatement parsedStatement = mock(ParsedStatement.class);
274+
when(parsedStatement.getType()).thenReturn(StatementType.QUERY);
275+
when(parsedStatement.isQuery()).thenReturn(true);
276+
Statement statement = Statement.of("SELECT * FROM FOO");
277+
when(parsedStatement.getStatement()).thenReturn(statement);
278+
279+
ReadWriteTransaction transaction = createSubject();
280+
ResultSet rs =
281+
get(transaction.executeQueryAsync(CallType.SYNC, parsedStatement, AnalyzeMode.WITH_STATS));
282+
assertThat(rs, is(notNullValue()));
283+
while (rs.next()) {
284+
// do nothing
285+
}
286+
assertThat(rs.getStats(), is(notNullValue()));
287+
}
288+
289+
@Test
290+
public void testWithPlanAndStatsQuery() {
291+
ParsedStatement parsedStatement = mock(ParsedStatement.class);
292+
when(parsedStatement.getType()).thenReturn(StatementType.QUERY);
293+
when(parsedStatement.isQuery()).thenReturn(true);
294+
Statement statement = Statement.of("SELECT * FROM FOO");
295+
when(parsedStatement.getStatement()).thenReturn(statement);
296+
297+
ReadWriteTransaction transaction = createSubject();
298+
ResultSet rs =
299+
get(transaction.executeQueryAsync(CallType.SYNC, parsedStatement, AnalyzeMode.WITH_STATS));
300+
assertThat(rs, is(notNullValue()));
301+
while (rs.next()) {
302+
// do nothing
303+
}
304+
assertThat(rs.getStats(), is(notNullValue()));
305+
}
306+
271307
@Test
272308
public void testExecuteUpdate() {
273309
ParsedStatement parsedStatement = mock(ParsedStatement.class);

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITReadWriteAutocommitSpannerTest.java‎

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -183,7 +183,8 @@ public void test05_BatchUpdateWithException() {
183183
@Test
184184
public void test06_AnalyzeUpdate() {
185185
assumeFalse(
186-
"Emulator does not support PLAN and PROFILE", EmulatorSpannerHelper.isUsingEmulator());
186+
"Emulator does not support PLAN, PROFILE, WITH_STATS AND WITH_PLAN_AND_STATS",
187+
EmulatorSpannerHelper.isUsingEmulator());
187188

188189
// PLAN should not execute the update.
189190
try (ITConnection connection = createConnection()) {
@@ -216,5 +217,29 @@ public void test06_AnalyzeUpdate() {
216217
assertTrue(resultSetStats.hasRowCountExact());
217218
assertTrue(resultSetStats.getRowCountExact() > 0);
218219
}
220+
221+
try (ITConnection connection = createConnection()) {
222+
ResultSetStats resultSetStats =
223+
connection.analyzeUpdate(
224+
Statement.of("UPDATE TEST SET NAME='test_updated' WHERE ID > 0"),
225+
QueryAnalyzeMode.WITH_STATS);
226+
227+
// Executing the update in WITH_STATS mode should execute the update
228+
assertNotNull(resultSetStats);
229+
assertFalse(resultSetStats.hasQueryPlan());
230+
assertTrue(resultSetStats.hasQueryStats());
231+
}
232+
233+
try (ITConnection connection = createConnection()) {
234+
ResultSetStats resultSetStats =
235+
connection.analyzeUpdate(
236+
Statement.of("UPDATE TEST SET NAME='test_updated' WHERE ID > 0"),
237+
QueryAnalyzeMode.WITH_STATS);
238+
239+
// Executing the update in WITH_PLAN_AND_STATS mode should execute the update
240+
assertNotNull(resultSetStats);
241+
assertTrue(resultSetStats.hasQueryPlan());
242+
assertTrue(resultSetStats.hasQueryStats());
243+
}
219244
}
220245
}

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITQueryTest.java‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1389,6 +1389,58 @@ public void analyzeProfile() {
13891389
assertThat(receivedStats.hasQueryStats()).isTrue();
13901390
}
13911391

1392+
@Test
1393+
public void analyzeWithStats() {
1394+
assumeFalse("Emulator does not support Analyze WithStats", isUsingEmulator());
1395+
1396+
String query = "SELECT 1 AS data UNION ALL SELECT 2 AS data ORDER BY data";
1397+
if (dialect.dialect == Dialect.POSTGRESQL) {
1398+
// "Statements with set operations and ORDER BY are not supported"
1399+
query = "SELECT 1 AS data UNION ALL SELECT 2 AS data";
1400+
}
1401+
Statement statement = Statement.of(query);
1402+
ResultSet resultSet =
1403+
statement.analyzeQuery(
1404+
getClient(dialect.dialect).singleUse(TimestampBound.strong()),
1405+
QueryAnalyzeMode.WITH_STATS);
1406+
assertThat(resultSet.next()).isTrue();
1407+
assertThat(resultSet.getType()).isEqualTo(Type.struct(StructField.of("data", Type.int64())));
1408+
assertThat(resultSet.getLong(0)).isEqualTo(1);
1409+
assertThat(resultSet.next()).isTrue();
1410+
assertThat(resultSet.getLong(0)).isEqualTo(2);
1411+
assertThat(resultSet.next()).isFalse();
1412+
ResultSetStats receivedStats = resultSet.getStats();
1413+
assertThat(receivedStats).isNotNull();
1414+
assertThat(receivedStats.hasQueryPlan()).isFalse();
1415+
assertThat(receivedStats.hasQueryStats()).isTrue();
1416+
}
1417+
1418+
@Test
1419+
public void analyzeWithPlanAndStats() {
1420+
assumeFalse("Emulator does not support Analyze WithPlanAndStats", isUsingEmulator());
1421+
1422+
String query = "SELECT 1 AS data UNION ALL SELECT 2 AS data ORDER BY data";
1423+
if (dialect.dialect == Dialect.POSTGRESQL) {
1424+
// "Statements with set operations and ORDER BY are not supported"
1425+
query = "SELECT 1 AS data UNION ALL SELECT 2 AS data";
1426+
}
1427+
Statement statement = Statement.of(query);
1428+
ResultSet resultSet =
1429+
statement.analyzeQuery(
1430+
getClient(dialect.dialect).singleUse(TimestampBound.strong()),
1431+
QueryAnalyzeMode.WITH_PLAN_AND_STATS);
1432+
assertThat(resultSet.next()).isTrue();
1433+
assertThat(resultSet.getType()).isEqualTo(Type.struct(StructField.of("data", Type.int64())));
1434+
assertThat(resultSet.getLong(0)).isEqualTo(1);
1435+
assertThat(resultSet.next()).isTrue();
1436+
assertThat(resultSet.getLong(0)).isEqualTo(2);
1437+
assertThat(resultSet.next()).isFalse();
1438+
ResultSetStats receivedStats = resultSet.getStats();
1439+
assertThat(receivedStats).isNotNull();
1440+
assertThat(receivedStats.hasQueryPlan()).isTrue();
1441+
assertThat(receivedStats.hasQueryStats()).isTrue();
1442+
}
1443+
13921444
@Test
13931445
public void testSelectArrayOfStructs() {
13941446
assumeFalse("structs are not supported on POSTGRESQL", dialect.dialect == Dialect.POSTGRESQL);

0 commit comments

Comments
 (0)