Skip to content

Commit d4f4546

Browse files
keshavdandevakeshavdandeva
committed
feat(bigquery): expose SessionInfo on TableResult
1 parent 6003ced commit d4f4546

5 files changed

Lines changed: 58 additions & 14 deletions

File tree

‎java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@
4242
import com.google.cloud.bigquery.BigQueryRetryHelper.BigQueryRetryHelperException;
4343
import com.google.cloud.bigquery.InsertAllRequest.RowToInsert;
4444
import com.google.cloud.bigquery.JobStatistics.QueryStatistics.StatementType;
45+
import com.google.cloud.bigquery.JobStatistics.SessionInfo;
4546
import com.google.cloud.bigquery.spi.v2.BigQueryRpc;
4647
import com.google.cloud.bigquery.spi.v2.HttpBigQueryRpc;
4748
import com.google.common.annotations.VisibleForTesting;
@@ -2104,6 +2105,8 @@ public com.google.api.services.bigquery.model.QueryResponse call()
21042105
Long totalBytesProcessed = results.getTotalBytesProcessed();
21052106
Long totalSlotMs = results.getTotalSlotMs();
21062107
Long numDmlAffectedRows = results.getNumDmlAffectedRows();
2108+
SessionInfo sessionInfo =
2109+
results.getSessionInfo() != null ? SessionInfo.fromPb(results.getSessionInfo()) : null;
21072110

21082111
if (results.getPageToken() != null) {
21092112
JobId jobId = JobId.fromPb(results.getJobReference());
@@ -2129,6 +2132,7 @@ public com.google.api.services.bigquery.model.QueryResponse call()
21292132
.setTotalBytesProcessed(totalBytesProcessed)
21302133
.setTotalSlotMs(totalSlotMs)
21312134
.setNumDmlAffectedRows(numDmlAffectedRows)
2135+
.setSessionInfo(sessionInfo)
21322136
.build();
21332137
}
21342138
// only 1 page of result
@@ -2154,6 +2158,7 @@ public com.google.api.services.bigquery.model.QueryResponse call()
21542158
.setTotalBytesProcessed(totalBytesProcessed)
21552159
.setTotalSlotMs(totalSlotMs)
21562160
.setNumDmlAffectedRows(numDmlAffectedRows)
2161+
.setSessionInfo(sessionInfo)
21572162
.build();
21582163
}
21592164

‎java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Job.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
import com.google.cloud.bigquery.JobConfiguration.Type;
3131
import com.google.cloud.bigquery.JobStatistics.QueryStatistics;
3232
import com.google.cloud.bigquery.JobStatistics.QueryStatistics.StatementType;
33+
import com.google.cloud.bigquery.JobStatistics.SessionInfo;
3334
import com.google.common.collect.ImmutableList;
3435
import io.opentelemetry.api.common.Attributes;
3536
import io.opentelemetry.api.trace.Span;
@@ -425,6 +426,7 @@ public TableResult getQueryResults(QueryResultsOption... options)
425426
Long totalBytesProcessed = stats != null ? stats.getTotalBytesProcessed() : null;
426427
Long totalSlotMs = stats != null ? stats.getTotalSlotMs() : null;
427428
Long numDmlAffectedRows = stats != null ? stats.getNumDmlAffectedRows() : null;
429+
SessionInfo sessionInfo = stats != null ? stats.getSessionInfo() : null;
428430

429431
// If there are no rows in the result, this may have been a DDL query.
430432
// Listing table data might fail, such as with CREATE VIEW queries.
@@ -442,6 +444,7 @@ public TableResult getQueryResults(QueryResultsOption... options)
442444
.setTotalBytesProcessed(totalBytesProcessed)
443445
.setTotalSlotMs(totalSlotMs)
444446
.setNumDmlAffectedRows(numDmlAffectedRows)
447+
.setSessionInfo(sessionInfo)
445448
.build();
446449
return emptyTableResult;
447450
}
@@ -461,6 +464,7 @@ public TableResult getQueryResults(QueryResultsOption... options)
461464
.setTotalBytesProcessed(totalBytesProcessed)
462465
.setTotalSlotMs(totalSlotMs)
463466
.setNumDmlAffectedRows(numDmlAffectedRows)
467+
.setSessionInfo(sessionInfo)
464468
.build();
465469
return tableResultWithJobId;
466470
} finally {

‎java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableResult.java‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import com.google.api.gax.paging.Page;
2020
import com.google.auto.value.AutoValue;
2121
import com.google.cloud.bigquery.JobStatistics.QueryStatistics.StatementType;
22+
import com.google.cloud.bigquery.JobStatistics.SessionInfo;
2223
import com.google.common.base.Function;
2324
import com.google.common.base.MoreObjects;
2425
import com.google.common.collect.Iterables;
@@ -62,6 +63,8 @@ public abstract static class Builder {
6263

6364
public abstract TableResult.Builder setNumDmlAffectedRows(Long numDmlAffectedRows);
6465

66+
public abstract TableResult.Builder setSessionInfo(SessionInfo sessionInfo);
67+
6568
/** Creates a @code TableResult} object. */
6669
public abstract TableResult build();
6770
}
@@ -124,6 +127,10 @@ public static Builder newBuilder() {
124127
@Nullable
125128
public abstract Long getNumDmlAffectedRows();
126129

130+
/** Returns information about the session if this query is part of one, if available. */
131+
@Nullable
132+
public abstract SessionInfo getSessionInfo();
133+
127134
@Override
128135
public boolean hasNextPage() {
129136
return getPageNoSchema().hasNextPage();
@@ -151,6 +158,7 @@ public TableResult getNextPage() {
151158
.setTotalBytesProcessed(getTotalBytesProcessed())
152159
.setTotalSlotMs(getTotalSlotMs())
153160
.setNumDmlAffectedRows(getNumDmlAffectedRows())
161+
.setSessionInfo(getSessionInfo())
154162
.build();
155163
}
156164
return null;
@@ -194,6 +202,7 @@ public String toString() {
194202
.add("totalBytesProcessed", getTotalBytesProcessed())
195203
.add("totalSlotMs", getTotalSlotMs())
196204
.add("numDmlAffectedRows", getNumDmlAffectedRows())
205+
.add("sessionInfo", getSessionInfo())
197206
.toString();
198207
}
199208

@@ -209,7 +218,8 @@ public final int hashCode() {
209218
getTotalBytesBilled(),
210219
getTotalBytesProcessed(),
211220
getTotalSlotMs(),
212-
getNumDmlAffectedRows());
221+
getNumDmlAffectedRows(),
222+
getSessionInfo());
213223
}
214224

215225
@Override
@@ -231,6 +241,7 @@ && getTotalRows() == response.getTotalRows()
231241
&& Objects.equals(getTotalBytesBilled(), response.getTotalBytesBilled())
232242
&& Objects.equals(getTotalBytesProcessed(), response.getTotalBytesProcessed())
233243
&& Objects.equals(getTotalSlotMs(), response.getTotalSlotMs())
234-
&& Objects.equals(getNumDmlAffectedRows(), response.getNumDmlAffectedRows());
244+
&& Objects.equals(getNumDmlAffectedRows(), response.getNumDmlAffectedRows())
245+
&& Objects.equals(getSessionInfo(), response.getSessionInfo());
235246
}
236247
}

‎java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@
5252
import com.google.api.services.bigquery.model.ProjectList;
5353
import com.google.api.services.bigquery.model.ProjectReference;
5454
import com.google.api.services.bigquery.model.QueryRequest;
55+
import com.google.api.services.bigquery.model.SessionInfo;
5556
import com.google.api.services.bigquery.model.TableCell;
5657
import com.google.api.services.bigquery.model.TableDataInsertAllRequest;
5758
import com.google.api.services.bigquery.model.TableDataInsertAllResponse;
@@ -170,6 +171,8 @@ public class BigQueryImplTest {
170171
.setField("timestampField");
171172
private static final TimePartitioning TIME_PARTITIONING_NULL_TYPE =
172173
TimePartitioning.fromPb(PB_TIMEPARTITIONING);
174+
private static final String SESSION_ID = "test-session-id";
175+
private static final SessionInfo PB_SESSION_INFO = new SessionInfo().setSessionId(SESSION_ID);
173176
private static final ImmutableMap<String, String> LABELS = ImmutableMap.of("key", "value");
174177
private static final StandardTableDefinition TABLE_DEFINITION_WITH_PARTITIONING =
175178
StandardTableDefinition.newBuilder()
@@ -2880,6 +2883,7 @@ void testQueryWithTimeoutSetsTimeout() throws InterruptedException, IOException
28802883
.setTotalBytesProcessed(42L)
28812884
.setTotalSlotMs(50L)
28822885
.setNumDmlAffectedRows(0L)
2886+
.setSessionInfo(PB_SESSION_INFO)
28832887
.setTotalRows(BigInteger.valueOf(1L));
28842888

28852889
when(bigqueryRpcMock.queryRpcSkipExceptionTranslation(eq(PROJECT), requestPbCapture.capture()))
@@ -2894,6 +2898,8 @@ void testQueryWithTimeoutSetsTimeout() throws InterruptedException, IOException
28942898
assertEquals((Long) 42L, tableResult.getTotalBytesProcessed());
28952899
assertEquals((Long) 50L, tableResult.getTotalSlotMs());
28962900
assertEquals((Long) 0L, tableResult.getNumDmlAffectedRows());
2901+
assertNotNull(tableResult.getSessionInfo());
2902+
assertEquals(SESSION_ID, tableResult.getSessionInfo().getSessionId());
28972903
QueryRequest requestPb = requestPbCapture.getValue();
28982904
assertEquals((Long) 1000L, requestPb.getTimeoutMs());
28992905
}

‎java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/TableResultTest.java‎

Lines changed: 30 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@
2121

2222
import com.google.api.gax.paging.Page;
2323
import com.google.cloud.PageImpl;
24+
import com.google.cloud.bigquery.JobStatistics.QueryStatistics.StatementType;
25+
import com.google.cloud.bigquery.JobStatistics.SessionInfo;
2426
import com.google.common.collect.ImmutableList;
2527
import org.junit.jupiter.api.Test;
2628

@@ -46,6 +48,15 @@ public Page<FieldValueList> getNextPage() {
4648
null,
4749
ImmutableList.of(newFieldValueList("2")));
4850
private static final Schema SCHEMA = Schema.of(Field.of("field", LegacySQLTypeName.INTEGER));
51+
private static final String SESSION_ID = "session_123";
52+
private static final SessionInfo SESSION_INFO =
53+
SessionInfo.newBuilder().setSessionId(SESSION_ID).build();
54+
private static final String SESSION_ID_1 = "session_1";
55+
private static final SessionInfo SESSION_INFO_1 =
56+
SessionInfo.newBuilder().setSessionId(SESSION_ID_1).build();
57+
private static final String SESSION_ID_2 = "session_2";
58+
private static final SessionInfo SESSION_INFO_2 =
59+
SessionInfo.newBuilder().setSessionId(SESSION_ID_2).build();
4960

5061
private static FieldValueList newFieldValueList(String s) {
5162
return FieldValueList.of(ImmutableList.of(FieldValue.of(PRIMITIVE, s)));
@@ -124,27 +135,29 @@ void testStatementTypeAndExecutionStats() {
124135
.setTotalRows(3L)
125136
.setPageNoSchema(INNER_PAGE_0)
126137
.setRowsInPage(2L)
127-
.setStatementType(JobStatistics.QueryStatistics.StatementType.SELECT)
138+
.setStatementType(StatementType.SELECT)
128139
.setTotalBytesBilled(1024L)
129140
.setTotalBytesProcessed(2048L)
130141
.setTotalSlotMs(500L)
131142
.setNumDmlAffectedRows(0L)
143+
.setSessionInfo(SESSION_INFO)
132144
.build();
133145

134-
assertThat(result.getStatementType())
135-
.isEqualTo(JobStatistics.QueryStatistics.StatementType.SELECT);
146+
assertThat(result.getStatementType()).isEqualTo(StatementType.SELECT);
136147
assertThat(result.getTotalBytesBilled()).isEqualTo(1024L);
137148
assertThat(result.getTotalBytesProcessed()).isEqualTo(2048L);
138149
assertThat(result.getTotalSlotMs()).isEqualTo(500L);
139150
assertThat(result.getNumDmlAffectedRows()).isEqualTo(0L);
151+
assertThat(result.getSessionInfo()).isEqualTo(SESSION_INFO);
152+
assertThat(result.getSessionInfo().getSessionId()).isEqualTo(SESSION_ID);
140153

141154
TableResult next = result.getNextPage();
142-
assertThat(next.getStatementType())
143-
.isEqualTo(JobStatistics.QueryStatistics.StatementType.SELECT);
155+
assertThat(next.getStatementType()).isEqualTo(StatementType.SELECT);
144156
assertThat(next.getTotalBytesBilled()).isEqualTo(1024L);
145157
assertThat(next.getTotalBytesProcessed()).isEqualTo(2048L);
146158
assertThat(next.getTotalSlotMs()).isEqualTo(500L);
147159
assertThat(next.getNumDmlAffectedRows()).isEqualTo(0L);
160+
assertThat(next.getSessionInfo()).isEqualTo(SESSION_INFO);
148161
}
149162

150163
@Test
@@ -155,23 +168,24 @@ void testToBuilder() {
155168
.setTotalRows(3L)
156169
.setPageNoSchema(INNER_PAGE_0)
157170
.setRowsInPage(2L)
158-
.setStatementType(JobStatistics.QueryStatistics.StatementType.INSERT)
171+
.setStatementType(StatementType.INSERT)
159172
.setTotalBytesBilled(500L)
160173
.setTotalBytesProcessed(1000L)
161174
.setTotalSlotMs(250L)
162175
.setNumDmlAffectedRows(5L)
176+
.setSessionInfo(SESSION_INFO)
163177
.build();
164178

165179
TableResult modified =
166180
result.toBuilder()
167-
.setStatementType(JobStatistics.QueryStatistics.StatementType.UPDATE)
181+
.setStatementType(StatementType.UPDATE)
168182
.setNumDmlAffectedRows(10L)
169183
.build();
170184

171-
assertThat(modified.getStatementType())
172-
.isEqualTo(JobStatistics.QueryStatistics.StatementType.UPDATE);
185+
assertThat(modified.getStatementType()).isEqualTo(StatementType.UPDATE);
173186
assertThat(modified.getNumDmlAffectedRows()).isEqualTo(10L);
174187
assertThat(modified.getTotalBytesBilled()).isEqualTo(500L);
188+
assertThat(modified.getSessionInfo()).isEqualTo(SESSION_INFO);
175189
}
176190

177191
@Test
@@ -182,11 +196,12 @@ void testEqualsAndHashCode() {
182196
.setTotalRows(3L)
183197
.setPageNoSchema(INNER_PAGE_0)
184198
.setRowsInPage(2L)
185-
.setStatementType(JobStatistics.QueryStatistics.StatementType.SELECT)
199+
.setStatementType(StatementType.SELECT)
186200
.setTotalBytesBilled(100L)
187201
.setTotalBytesProcessed(200L)
188202
.setTotalSlotMs(50L)
189203
.setNumDmlAffectedRows(0L)
204+
.setSessionInfo(SESSION_INFO_1)
190205
.build();
191206

192207
TableResult result2 =
@@ -195,11 +210,12 @@ void testEqualsAndHashCode() {
195210
.setTotalRows(3L)
196211
.setPageNoSchema(INNER_PAGE_0)
197212
.setRowsInPage(2L)
198-
.setStatementType(JobStatistics.QueryStatistics.StatementType.SELECT)
213+
.setStatementType(StatementType.SELECT)
199214
.setTotalBytesBilled(100L)
200215
.setTotalBytesProcessed(200L)
201216
.setTotalSlotMs(50L)
202217
.setNumDmlAffectedRows(0L)
218+
.setSessionInfo(SESSION_INFO_1)
203219
.build();
204220

205221
TableResult result3 =
@@ -208,17 +224,19 @@ void testEqualsAndHashCode() {
208224
.setTotalRows(3L)
209225
.setPageNoSchema(INNER_PAGE_0)
210226
.setRowsInPage(2L)
211-
.setStatementType(JobStatistics.QueryStatistics.StatementType.DELETE)
227+
.setStatementType(StatementType.DELETE)
212228
.setTotalBytesBilled(100L)
213229
.setTotalBytesProcessed(200L)
214230
.setTotalSlotMs(50L)
215231
.setNumDmlAffectedRows(1L)
232+
.setSessionInfo(SESSION_INFO_2)
216233
.build();
217234

218235
assertThat(result1).isEqualTo(result2);
219236
assertThat(result1.hashCode()).isEqualTo(result2.hashCode());
220237
assertThat(result1).isNotEqualTo(result3);
221238
assertThat(result1.toString()).contains("statementType=SELECT");
222239
assertThat(result1.toString()).contains("totalBytesBilled=100");
240+
assertThat(result1.toString()).contains("sessionId=" + SESSION_ID_1);
223241
}
224242
}

0 commit comments

Comments
 (0)