From 465ab61bc697ec0a946ea3a0b4864a3486fdbda3 Mon Sep 17 00:00:00 2001 From: Naresh Chaudhary Date: Mon, 3 Mar 2025 21:12:12 +0000 Subject: [PATCH 1/3] Add new QueryMode's support in client library --- .../cloud/spanner/AbstractReadContext.java | 6 ++++ .../com/google/cloud/spanner/ReadContext.java | 12 ++++++- .../com/google/cloud/spanner/ResultSet.java | 5 +-- .../cloud/spanner/TransactionRunnerImpl.java | 6 ++++ .../cloud/spanner/connection/AnalyzeMode.java | 25 +++++++++++--- .../cloud/spanner/DatabaseClientImplTest.java | 34 +++++++++++++++++++ 6 files changed, 81 insertions(+), 7 deletions(-) diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java index 145ad67f827..4df9123fa0d 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java @@ -618,6 +618,12 @@ public final ResultSet analyzeQuery(Statement statement, QueryAnalyzeMode readCo case PLAN: return executeQueryInternal( statement, com.google.spanner.v1.ExecuteSqlRequest.QueryMode.PLAN); + case WITH_STATS: + return executeQueryInternal( + statement, com.google.spanner.v1.ExecuteSqlRequest.QueryMode.WITH_STATS); + case WITH_PLAN_AND_STATS: + return executeQueryInternal( + statement, com.google.spanner.v1.ExecuteSqlRequest.QueryMode.WITH_PLAN_AND_STATS); default: throw new IllegalStateException( "Unknown value for QueryAnalyzeMode : " + readContextQueryMode); diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ReadContext.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ReadContext.java index c5dddfe1159..54a5d8b630f 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ReadContext.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ReadContext.java @@ -34,7 +34,17 @@ enum QueryAnalyzeMode { /** Retrieves only the query plan information. No result data is returned. */ PLAN, /** Retrieves both query plan and query execution statistics along with the result data. */ - PROFILE + PROFILE, + /** + * Retrieves the overall (but not operator-level) execution statistics along with the result + * data. + */ + WITH_STATS, + /** + * Retrieves the query plan, overall (but not operator-level) execution statistics along with + * the result data. + */ + WITH_PLAN_AND_STATS } /** * Reads zero or more rows from a database. diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResultSet.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResultSet.java index cd6fa10b996..6b219a76a42 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResultSet.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResultSet.java @@ -63,8 +63,9 @@ public interface ResultSet extends AutoCloseable, StructReader { void close(); /** - * Returns the {@link ResultSetStats} for the query only if the query was executed in either the - * {@code PLAN} or the {@code PROFILE} mode via the {@link ReadContext#analyzeQuery(Statement, + * Returns the {@link ResultSetStats} for the query only if the query was executed in {@code + * PLAN}, {@code PROFILE}, {@code WITH_STATS} or the {@code WITH_PLAN_AND_STATS} mode via the + * {@link ReadContext#analyzeQuery(Statement, * com.google.cloud.spanner.ReadContext.QueryAnalyzeMode)} method or for DML statements in {@link * ReadContext#executeQuery(Statement, QueryOption...)}. Attempts to call this method on a {@code * ResultSet} not obtained from {@code analyzeQuery} or {@code executeQuery} will return a {@code diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java index fad4ce564ab..b9155772361 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java @@ -878,6 +878,12 @@ private ResultSet internalAnalyzeStatement( case PROFILE: queryMode = QueryMode.PROFILE; break; + case WITH_STATS: + queryMode = QueryMode.WITH_STATS; + break; + case WITH_PLAN_AND_STATS: + queryMode = QueryMode.WITH_PLAN_AND_STATS; + break; default: throw SpannerExceptionFactory.newSpannerException( ErrorCode.INVALID_ARGUMENT, "Unknown analyze mode: " + analyzeMode); diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AnalyzeMode.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AnalyzeMode.java index f67d2267771..191621b3e70 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AnalyzeMode.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AnalyzeMode.java @@ -19,14 +19,27 @@ import com.google.cloud.spanner.ReadContext.QueryAnalyzeMode; /** - * {@link AnalyzeMode} indicates whether a query should be executed as a normal query (NONE), - * whether only a query plan should be returned, or whether the query should be profiled while - * executed. + * {@link AnalyzeMode} controls the execution and returned information for a query: + * + * */ enum AnalyzeMode { NONE(null), PLAN(QueryAnalyzeMode.PLAN), - PROFILE(QueryAnalyzeMode.PROFILE); + PROFILE(QueryAnalyzeMode.PROFILE), + WITH_STATS(QueryAnalyzeMode.WITH_STATS), + WITH_PLAN_AND_STATS(QueryAnalyzeMode.WITH_PLAN_AND_STATS); private final QueryAnalyzeMode mode; @@ -45,6 +58,10 @@ static AnalyzeMode of(QueryAnalyzeMode mode) { return AnalyzeMode.PLAN; case PROFILE: return AnalyzeMode.PROFILE; + case WITH_STATS: + return AnalyzeMode.WITH_STATS; + case WITH_PLAN_AND_STATS: + return AnalyzeMode.WITH_PLAN_AND_STATS; default: throw new IllegalArgumentException(mode + " is unknown"); } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java index 87ea5c19ce9..ba35eedc37f 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java @@ -3538,6 +3538,40 @@ public void testBackendQueryOptionsWithAnalyzeQuery() { } } + @Test + public void testWithStatsQueryModeWithAnalyzeQuery() { + // Use a Spanner instance with MinSession=0 to prevent background requests + // from the session pool interfering with the test case. + try (Spanner spanner = + SpannerOptions.newBuilder() + .setProjectId("[PROJECT]") + .setChannelProvider(channelProvider) + .setCredentials(NoCredentials.getInstance()) + .setSessionPoolOption(SessionPoolOptions.newBuilder().setMinSessions(0).build()) + .build() + .getService()) { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE")); + try (ReadOnlyTransaction tx = client.readOnlyTransaction()) { + try (ResultSet rs = + tx.analyzeQuery( + Statement.newBuilder(SELECT1.getSql()) + .build(), + QueryAnalyzeMode.WITH_STATS)) { + // Just iterate over the results to execute the query. + consumeResults(rs); + } + } + // Check that the last query was executed using a custom optimizer version and statistics + // package. + List requests = mockSpanner.getRequests(); + assertThat(requests).isNotEmpty(); + assertThat(requests.get(requests.size() - 1)).isInstanceOf(ExecuteSqlRequest.class); + ExecuteSqlRequest request = (ExecuteSqlRequest) requests.get(requests.size() - 1); + assertThat(request.getQueryMode()).isEqualTo(QueryMode.WITH_STATS); + } + } + @Test public void testBackendPartitionQueryOptions() { // Use a Spanner instance with MinSession=0 to prevent background requests From b6f2949708ba718652b86f365e7fd690f5c3bb58 Mon Sep 17 00:00:00 2001 From: Naresh Chaudhary Date: Tue, 4 Mar 2025 08:35:26 +0000 Subject: [PATCH 2/3] Add tests for new queryMode's --- .../cloud/spanner/connection/Connection.java | 14 ++++- .../cloud/spanner/DatabaseClientImplTest.java | 4 +- .../cloud/spanner/GrpcResultSetTest.java | 57 +++++++++++++++++++ .../connection/ConnectionImplTest.java | 14 +++++ .../connection/ReadOnlyTransactionTest.java | 48 ++++++++++++++++ .../connection/ReadWriteTransactionTest.java | 36 ++++++++++++ .../it/ITReadWriteAutocommitSpannerTest.java | 27 ++++++++- .../google/cloud/spanner/it/ITQueryTest.java | 52 +++++++++++++++++ 8 files changed, 245 insertions(+), 7 deletions(-) diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/Connection.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/Connection.java index 7bf4e47bd9a..147d8df1584 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/Connection.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/Connection.java @@ -1327,9 +1327,13 @@ PartitionedQueryResultSet runPartitionedQuery( * Analyzes a DML statement and returns query plan and/or execution statistics information. * *

{@link com.google.cloud.spanner.ReadContext.QueryAnalyzeMode#PLAN} only returns the plan for - * the statement. {@link com.google.cloud.spanner.ReadContext.QueryAnalyzeMode#PROFILE} executes - * the DML statement, returns the modified row count and execution statistics, and the effects of - * the DML statement will be visible to subsequent operations in the transaction. + * the statement. {@link com.google.cloud.spanner.ReadContext.QueryAnalyzeMode#WITH_STATS} returns + * the overall (but not operator-level) execution statistics. {@link + * com.google.cloud.spanner.ReadContext.QueryAnalyzeMode#WITH_PLAN_AND_STATS} returns the query + * plan and overall (but not operator-level) execution statistics. {@link + * com.google.cloud.spanner.ReadContext.QueryAnalyzeMode#PROFILE} executes the DML statement, + * returns the modified row count and execution statistics, and the effects of the DML statement + * will be visible to subsequent operations in the transaction. * * @deprecated Use {@link #analyzeUpdateStatement(Statement, QueryAnalyzeMode, UpdateOption...)} * instead @@ -1345,6 +1349,10 @@ default ResultSetStats analyzeUpdate(Statement update, QueryAnalyzeMode analyzeM * *

{@link com.google.cloud.spanner.ReadContext.QueryAnalyzeMode#PLAN} only returns the plan and * undeclared parameters for the statement. {@link + * com.google.cloud.spanner.ReadContext.QueryAnalyzeMode#WITH_STATS} returns the overall (but not + * operator-level) execution statistics. {@link + * com.google.cloud.spanner.ReadContext.QueryAnalyzeMode#WITH_PLAN_AND_STATS} returns the query + * plan and overall (but not operator-level) execution statistics. {@link * com.google.cloud.spanner.ReadContext.QueryAnalyzeMode#PROFILE} also executes the DML statement, * returns the modified row count and execution statistics, and the effects of the DML statement * will be visible to subsequent operations in the transaction. diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java index ba35eedc37f..afc4713372e 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java @@ -3555,9 +3555,7 @@ public void testWithStatsQueryModeWithAnalyzeQuery() { try (ReadOnlyTransaction tx = client.readOnlyTransaction()) { try (ResultSet rs = tx.analyzeQuery( - Statement.newBuilder(SELECT1.getSql()) - .build(), - QueryAnalyzeMode.WITH_STATS)) { + Statement.newBuilder(SELECT1.getSql()).build(), QueryAnalyzeMode.WITH_STATS)) { // Just iterate over the results to execute the query. consumeResults(rs); } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcResultSetTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcResultSetTest.java index 25c01560e92..a958c59dfee 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcResultSetTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcResultSetTest.java @@ -470,6 +470,63 @@ public void planResult() { resultSet.close(); } + @Test + public void withStatsResult() { + Map statsMap = + ImmutableMap.of( + "f1", Value.string("").toProto(), + "f2", Value.string("").toProto()); + ResultSetStats stats = + ResultSetStats.newBuilder() + .setQueryStats(com.google.protobuf.Struct.newBuilder().putAllFields(statsMap).build()) + .build(); + ArrayList dataType = new ArrayList<>(); + dataType.add(Type.StructField.of("data", Type.string())); + consumer.onPartialResultSet( + PartialResultSet.newBuilder() + .setMetadata(makeMetadata(Type.struct(dataType))) + .addValues(Value.string("d1").toProto()) + .setChunkedValue(false) + .setStats(stats) + .build()); + resultSet = resultSetWithMode(QueryMode.WITH_STATS); + consumer.onCompleted(); + assertThat(resultSet.next()).isTrue(); + assertThat(resultSet.next()).isFalse(); + ResultSetStats receivedStats = resultSet.getStats(); + assertThat(stats).isEqualTo(receivedStats); + resultSet.close(); + } + + @Test + public void withPlanAndStatsResult() { + Map statsMap = + ImmutableMap.of( + "f1", Value.string("").toProto(), + "f2", Value.string("").toProto()); + ResultSetStats stats = + ResultSetStats.newBuilder() + .setQueryPlan(QueryPlan.newBuilder().build()) + .setQueryStats(com.google.protobuf.Struct.newBuilder().putAllFields(statsMap).build()) + .build(); + ArrayList dataType = new ArrayList<>(); + dataType.add(Type.StructField.of("data", Type.string())); + consumer.onPartialResultSet( + PartialResultSet.newBuilder() + .setMetadata(makeMetadata(Type.struct(dataType))) + .addValues(Value.string("d1").toProto()) + .setChunkedValue(false) + .setStats(stats) + .build()); + resultSet = resultSetWithMode(QueryMode.WITH_PLAN_AND_STATS); + consumer.onCompleted(); + assertThat(resultSet.next()).isTrue(); + assertThat(resultSet.next()).isFalse(); + ResultSetStats receivedStats = resultSet.getStats(); + assertThat(stats).isEqualTo(receivedStats); + resultSet.close(); + } + @Test public void statsUnavailable() { ResultSetStats stats = ResultSetStats.newBuilder().build(); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionImplTest.java index d4b7c035658..94af2859d75 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionImplTest.java @@ -272,6 +272,11 @@ public static ConnectionImpl createConnection(final ConnectionOptions options, D .thenReturn(select1ResultSetWithStats); when(singleUseReadOnlyTx.analyzeQuery(Statement.of(SELECT), QueryAnalyzeMode.PROFILE)) .thenReturn(select1ResultSetWithStats); + when(singleUseReadOnlyTx.analyzeQuery(Statement.of(SELECT), QueryAnalyzeMode.WITH_STATS)) + .thenReturn(select1ResultSetWithStats); + when(singleUseReadOnlyTx.analyzeQuery( + Statement.of(SELECT), QueryAnalyzeMode.WITH_PLAN_AND_STATS)) + .thenReturn(select1ResultSetWithStats); when(singleUseReadOnlyTx.getReadTimestamp()) .then( invocation -> { @@ -307,6 +312,11 @@ public static ConnectionImpl createConnection(final ConnectionOptions options, D .thenReturn(select1ResultSetWithStats); when(txContext.analyzeQuery(Statement.of(SELECT), QueryAnalyzeMode.PROFILE)) .thenReturn(select1ResultSetWithStats); + when(txContext.analyzeQuery(Statement.of(SELECT), QueryAnalyzeMode.WITH_STATS)) + .thenReturn(select1ResultSetWithStats); + when(txContext.analyzeQuery( + Statement.of(SELECT), QueryAnalyzeMode.WITH_PLAN_AND_STATS)) + .thenReturn(select1ResultSetWithStats); when(txContext.executeUpdate(Statement.of(UPDATE))).thenReturn(1L); return new SimpleTransactionManager(txContext, options.isReturnCommitStats()); }); @@ -328,6 +338,10 @@ public static ConnectionImpl createConnection(final ConnectionOptions options, D .thenReturn(select1ResultSetWithStats); when(tx.analyzeQuery(Statement.of(SELECT), QueryAnalyzeMode.PROFILE)) .thenReturn(select1ResultSetWithStats); + when(tx.analyzeQuery(Statement.of(SELECT), QueryAnalyzeMode.WITH_STATS)) + .thenReturn(select1ResultSetWithStats); + when(tx.analyzeQuery(Statement.of(SELECT), QueryAnalyzeMode.WITH_PLAN_AND_STATS)) + .thenReturn(select1ResultSetWithStats); when(tx.getReadTimestamp()) .then( ignored -> { diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadOnlyTransactionTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadOnlyTransactionTest.java index e243fbd620a..023b1a86134 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadOnlyTransactionTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadOnlyTransactionTest.java @@ -358,6 +358,54 @@ public void testPlanQuery() { } } + @Test + public void testWithStatsQuery() { + for (TimestampBound staleness : getTestTimestampBounds()) { + ParsedStatement parsedStatement = mock(ParsedStatement.class); + when(parsedStatement.getType()).thenReturn(StatementType.QUERY); + when(parsedStatement.isQuery()).thenReturn(true); + Statement statement = Statement.of("SELECT * FROM FOO"); + when(parsedStatement.getStatement()).thenReturn(statement); + when(parsedStatement.getSqlWithoutComments()).thenReturn(statement.getSql()); + + ReadOnlyTransaction transaction = createSubject(staleness); + ResultSet rs = + get( + transaction.executeQueryAsync( + CallType.SYNC, parsedStatement, AnalyzeMode.WITH_STATS)); + assertThat(rs, is(notNullValue())); + // get all results and then get the stats + while (rs.next()) { + // do nothing + } + assertThat(rs.getStats(), is(notNullValue())); + } + } + + @Test + public void testWithPlanAndStatsQuery() { + for (TimestampBound staleness : getTestTimestampBounds()) { + ParsedStatement parsedStatement = mock(ParsedStatement.class); + when(parsedStatement.getType()).thenReturn(StatementType.QUERY); + when(parsedStatement.isQuery()).thenReturn(true); + Statement statement = Statement.of("SELECT * FROM FOO"); + when(parsedStatement.getStatement()).thenReturn(statement); + when(parsedStatement.getSqlWithoutComments()).thenReturn(statement.getSql()); + + ReadOnlyTransaction transaction = createSubject(staleness); + ResultSet rs = + get( + transaction.executeQueryAsync( + CallType.SYNC, parsedStatement, AnalyzeMode.WITH_PLAN_AND_STATS)); + assertThat(rs, is(notNullValue())); + // get all results and then get the stats + while (rs.next()) { + // do nothing + } + assertThat(rs.getStats(), is(notNullValue())); + } + } + @Test public void testProfileQuery() { for (TimestampBound staleness : getTestTimestampBounds()) { diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadWriteTransactionTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadWriteTransactionTest.java index 9fbb5b5bf16..87e010939d8 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadWriteTransactionTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadWriteTransactionTest.java @@ -268,6 +268,42 @@ public void testProfileQuery() { assertThat(rs.getStats(), is(notNullValue())); } + @Test + public void testWithStatsQuery() { + ParsedStatement parsedStatement = mock(ParsedStatement.class); + when(parsedStatement.getType()).thenReturn(StatementType.QUERY); + when(parsedStatement.isQuery()).thenReturn(true); + Statement statement = Statement.of("SELECT * FROM FOO"); + when(parsedStatement.getStatement()).thenReturn(statement); + + ReadWriteTransaction transaction = createSubject(); + ResultSet rs = + get(transaction.executeQueryAsync(CallType.SYNC, parsedStatement, AnalyzeMode.WITH_STATS)); + assertThat(rs, is(notNullValue())); + while (rs.next()) { + // do nothing + } + assertThat(rs.getStats(), is(notNullValue())); + } + + @Test + public void testWithPlanAndStatsQuery() { + ParsedStatement parsedStatement = mock(ParsedStatement.class); + when(parsedStatement.getType()).thenReturn(StatementType.QUERY); + when(parsedStatement.isQuery()).thenReturn(true); + Statement statement = Statement.of("SELECT * FROM FOO"); + when(parsedStatement.getStatement()).thenReturn(statement); + + ReadWriteTransaction transaction = createSubject(); + ResultSet rs = + get(transaction.executeQueryAsync(CallType.SYNC, parsedStatement, AnalyzeMode.WITH_STATS)); + assertThat(rs, is(notNullValue())); + while (rs.next()) { + // do nothing + } + assertThat(rs.getStats(), is(notNullValue())); + } + @Test public void testExecuteUpdate() { ParsedStatement parsedStatement = mock(ParsedStatement.class); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITReadWriteAutocommitSpannerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITReadWriteAutocommitSpannerTest.java index 8c9aaba823b..a795ac1c052 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITReadWriteAutocommitSpannerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITReadWriteAutocommitSpannerTest.java @@ -183,7 +183,8 @@ public void test05_BatchUpdateWithException() { @Test public void test06_AnalyzeUpdate() { assumeFalse( - "Emulator does not support PLAN and PROFILE", EmulatorSpannerHelper.isUsingEmulator()); + "Emulator does not support PLAN, PROFILE, WITH_STATS AND WITH_PLAN_AND_STATS", + EmulatorSpannerHelper.isUsingEmulator()); // PLAN should not execute the update. try (ITConnection connection = createConnection()) { @@ -216,5 +217,29 @@ public void test06_AnalyzeUpdate() { assertTrue(resultSetStats.hasRowCountExact()); assertTrue(resultSetStats.getRowCountExact() > 0); } + + try (ITConnection connection = createConnection()) { + ResultSetStats resultSetStats = + connection.analyzeUpdate( + Statement.of("UPDATE TEST SET NAME='test_updated' WHERE ID > 0"), + QueryAnalyzeMode.WITH_STATS); + + // Executing the update in WITH_STATS mode should execute the update + assertNotNull(resultSetStats); + assertFalse(resultSetStats.hasQueryPlan()); + assertTrue(resultSetStats.hasQueryStats()); + } + + try (ITConnection connection = createConnection()) { + ResultSetStats resultSetStats = + connection.analyzeUpdate( + Statement.of("UPDATE TEST SET NAME='test_updated' WHERE ID > 0"), + QueryAnalyzeMode.WITH_STATS); + + // Executing the update in WITH_PLAN_AND_STATS mode should execute the update + assertNotNull(resultSetStats); + assertTrue(resultSetStats.hasQueryPlan()); + assertTrue(resultSetStats.hasQueryStats()); + } } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITQueryTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITQueryTest.java index 18044c452b5..5ad69044c13 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITQueryTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITQueryTest.java @@ -1389,6 +1389,58 @@ public void analyzeProfile() { assertThat(receivedStats.hasQueryStats()).isTrue(); } + @Test + public void analyzeWithStats() { + assumeFalse("Emulator does not support Analyze WithStats", isUsingEmulator()); + + String query = "SELECT 1 AS data UNION ALL SELECT 2 AS data ORDER BY data"; + if (dialect.dialect == Dialect.POSTGRESQL) { + // "Statements with set operations and ORDER BY are not supported" + query = "SELECT 1 AS data UNION ALL SELECT 2 AS data"; + } + Statement statement = Statement.of(query); + ResultSet resultSet = + statement.analyzeQuery( + getClient(dialect.dialect).singleUse(TimestampBound.strong()), + QueryAnalyzeMode.WITH_STATS); + assertThat(resultSet.next()).isTrue(); + assertThat(resultSet.getType()).isEqualTo(Type.struct(StructField.of("data", Type.int64()))); + assertThat(resultSet.getLong(0)).isEqualTo(1); + assertThat(resultSet.next()).isTrue(); + assertThat(resultSet.getLong(0)).isEqualTo(2); + assertThat(resultSet.next()).isFalse(); + ResultSetStats receivedStats = resultSet.getStats(); + assertThat(receivedStats).isNotNull(); + assertThat(receivedStats.hasQueryPlan()).isFalse(); + assertThat(receivedStats.hasQueryStats()).isTrue(); + } + + @Test + public void analyzeWithPlanAndStats() { + assumeFalse("Emulator does not support Analyze WithPlanAndStats", isUsingEmulator()); + + String query = "SELECT 1 AS data UNION ALL SELECT 2 AS data ORDER BY data"; + if (dialect.dialect == Dialect.POSTGRESQL) { + // "Statements with set operations and ORDER BY are not supported" + query = "SELECT 1 AS data UNION ALL SELECT 2 AS data"; + } + Statement statement = Statement.of(query); + ResultSet resultSet = + statement.analyzeQuery( + getClient(dialect.dialect).singleUse(TimestampBound.strong()), + QueryAnalyzeMode.WITH_PLAN_AND_STATS); + assertThat(resultSet.next()).isTrue(); + assertThat(resultSet.getType()).isEqualTo(Type.struct(StructField.of("data", Type.int64()))); + assertThat(resultSet.getLong(0)).isEqualTo(1); + assertThat(resultSet.next()).isTrue(); + assertThat(resultSet.getLong(0)).isEqualTo(2); + assertThat(resultSet.next()).isFalse(); + ResultSetStats receivedStats = resultSet.getStats(); + assertThat(receivedStats).isNotNull(); + assertThat(receivedStats.hasQueryPlan()).isTrue(); + assertThat(receivedStats.hasQueryStats()).isTrue(); + } + @Test public void testSelectArrayOfStructs() { assumeFalse("structs are not supported on POSTGRESQL", dialect.dialect == Dialect.POSTGRESQL); From 33263f1fd24f0f859ef581ebfe0129a94a014621 Mon Sep 17 00:00:00 2001 From: Naresh Chaudhary Date: Thu, 13 Mar 2025 08:43:04 +0000 Subject: [PATCH 3/3] resolve comments --- .../src/main/java/com/google/cloud/spanner/ReadContext.java | 6 +++++- .../java/com/google/cloud/spanner/GrpcResultSetTest.java | 2 +- .../connection/it/ITReadWriteAutocommitSpannerTest.java | 2 +- 3 files changed, 7 insertions(+), 3 deletions(-) diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ReadContext.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ReadContext.java index 54a5d8b630f..d0b77364d48 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ReadContext.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ReadContext.java @@ -33,7 +33,11 @@ public interface ReadContext extends AutoCloseable { enum QueryAnalyzeMode { /** Retrieves only the query plan information. No result data is returned. */ PLAN, - /** Retrieves both query plan and query execution statistics along with the result data. */ + /** + * Retrieves the query plan, overall execution statistics, operator level execution statistics + * along with the result data. This has a performance overhead compared to the other modes. It + * isn't recommended to use this mode for production traffic. + */ PROFILE, /** * Retrieves the overall (but not operator-level) execution statistics along with the result diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcResultSetTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcResultSetTest.java index a958c59dfee..c6d47441ecd 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcResultSetTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcResultSetTest.java @@ -494,7 +494,7 @@ public void withStatsResult() { assertThat(resultSet.next()).isTrue(); assertThat(resultSet.next()).isFalse(); ResultSetStats receivedStats = resultSet.getStats(); - assertThat(stats).isEqualTo(receivedStats); + assertThat(receivedStats).isEqualTo(stats); resultSet.close(); } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITReadWriteAutocommitSpannerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITReadWriteAutocommitSpannerTest.java index a795ac1c052..74fe55400bc 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITReadWriteAutocommitSpannerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITReadWriteAutocommitSpannerTest.java @@ -183,7 +183,7 @@ public void test05_BatchUpdateWithException() { @Test public void test06_AnalyzeUpdate() { assumeFalse( - "Emulator does not support PLAN, PROFILE, WITH_STATS AND WITH_PLAN_AND_STATS", + "Emulator does not support PLAN, PROFILE, WITH_STATS, AND WITH_PLAN_AND_STATS", EmulatorSpannerHelper.isUsingEmulator()); // PLAN should not execute the update.