From acc6cb8ad318ae41a9a3a00a5942025c14cbe681 Mon Sep 17 00:00:00 2001 From: Stephanie Wang Date: Fri, 15 Oct 2021 15:09:25 -0400 Subject: [PATCH] feat: add session support (#1652) Adds: - CreateSession to Query config - SessionInfo in JobStatistics to record session stats --- .../google/cloud/bigquery/JobStatistics.java | 90 ++++++++++++++++++- .../cloud/bigquery/QueryJobConfiguration.java | 29 ++++++ .../cloud/bigquery/QueryRequestInfo.java | 7 ++ .../cloud/bigquery/JobStatisticsTest.java | 15 ++++ .../bigquery/QueryJobConfigurationTest.java | 3 + .../cloud/bigquery/QueryRequestInfoTest.java | 2 + .../cloud/bigquery/it/ITBigQueryTest.java | 33 +++++++ 7 files changed, 177 insertions(+), 2 deletions(-) diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java index a3d04c017..999f8d2fe 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java @@ -45,6 +45,7 @@ public abstract class JobStatistics implements Serializable { private final ScriptStatistics scriptStatistics; private final List reservationUsage; private final TransactionInfo transactionInfo; + private final SessionInfo sessionInfo; /** A Google BigQuery Copy Job statistics. */ public static class CopyStatistics extends JobStatistics { @@ -1251,6 +1252,76 @@ static TransactionInfo fromPb( } } + // SessionInfo contains information about the session if this job is part of one. + public static class SessionInfo { + + // Id of the session + private final String sessionId; + + public static class Builder { + + private String sessionId; + + private Builder() {}; + + Builder setSessionId(String sessionId) { + this.sessionId = sessionId; + return this; + } + + SessionInfo build() { + return new SessionInfo(this); + } + } + + private SessionInfo(Builder builder) { + this.sessionId = builder.sessionId; + } + + public String getSessionId() { + return sessionId; + } + + static Builder newBuilder() { + return new Builder(); + } + + ToStringHelper toStringHelper() { + return MoreObjects.toStringHelper(this).add("sessionId", sessionId); + } + + @Override + public String toString() { + return toStringHelper().toString(); + } + + @Override + public boolean equals(Object obj) { + return obj == this + || obj != null + && obj.getClass().equals(SessionInfo.class) + && Objects.equals(toPb(), ((SessionInfo) obj).toPb()); + } + + @Override + public int hashCode() { + return Objects.hash(sessionId); + } + + com.google.api.services.bigquery.model.SessionInfo toPb() { + com.google.api.services.bigquery.model.SessionInfo sessionInfo = + new com.google.api.services.bigquery.model.SessionInfo(); + sessionInfo.setSessionId(sessionId); + return sessionInfo; + } + + static SessionInfo fromPb(com.google.api.services.bigquery.model.SessionInfo sessionInfo) { + SessionInfo.Builder builder = newBuilder(); + builder.setSessionId(sessionInfo.getSessionId()); + return builder.build(); + } + } + abstract static class Builder> { private Long creationTime; @@ -1261,6 +1332,7 @@ abstract static class Builder> private ScriptStatistics scriptStatistics; private List reservationUsage; private TransactionInfo transactionInfo; + private SessionInfo sessionInfo; protected Builder() {} @@ -1280,6 +1352,9 @@ protected Builder(com.google.api.services.bigquery.model.JobStatistics statistic if (statisticsPb.getTransactionInfo() != null) { this.transactionInfo = TransactionInfo.fromPb(statisticsPb.getTransactionInfo()); } + if (statisticsPb.getSessionInfo() != null) { + this.sessionInfo = SessionInfo.fromPb(statisticsPb.getSessionInfo()); + } } @SuppressWarnings("unchecked") @@ -1314,6 +1389,7 @@ protected JobStatistics(Builder builder) { this.scriptStatistics = builder.scriptStatistics; this.reservationUsage = builder.reservationUsage; this.transactionInfo = builder.transactionInfo; + this.sessionInfo = builder.sessionInfo; } /** Returns the creation time of the job in milliseconds since epoch. */ @@ -1362,6 +1438,11 @@ public TransactionInfo getTransactionInfo() { return transactionInfo; } + /** Info of the session if this job is part of one. */ + public SessionInfo getSessionInfo() { + return sessionInfo; + } + ToStringHelper toStringHelper() { return MoreObjects.toStringHelper(this) .add("creationTime", creationTime) @@ -1371,7 +1452,8 @@ ToStringHelper toStringHelper() { .add("parentJobId", parentJobId) .add("scriptStatistics", scriptStatistics) .add("reservationUsage", reservationUsage) - .add("transactionInfo", transactionInfo); + .add("transactionInfo", transactionInfo) + .add("sessionInfo", sessionInfo); } @Override @@ -1388,7 +1470,8 @@ final int baseHashCode() { parentJobId, scriptStatistics, reservationUsage, - transactionInfo); + transactionInfo, + sessionInfo); } final boolean baseEquals(JobStatistics jobStatistics) { @@ -1413,6 +1496,9 @@ com.google.api.services.bigquery.model.JobStatistics toPb() { if (transactionInfo != null) { statistics.setTransactionInfo(transactionInfo.toPb()); } + if (sessionInfo != null) { + statistics.setSessionInfo(sessionInfo.toPb()); + } return statistics; } diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryJobConfiguration.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryJobConfiguration.java index 7733e48e6..48ec22caf 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryJobConfiguration.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryJobConfiguration.java @@ -56,6 +56,7 @@ public final class QueryJobConfiguration extends JobConfiguration { private final DatasetId defaultDataset; private final Priority priority; private final Boolean allowLargeResults; + private final Boolean createSession; private final Boolean useQueryCache; private final Boolean flattenResults; private final Boolean dryRun; @@ -108,6 +109,7 @@ public static final class Builder private DatasetId defaultDataset; private Priority priority; private Boolean allowLargeResults; + private Boolean createSession; private Boolean useQueryCache; private Boolean flattenResults; private Boolean dryRun; @@ -142,6 +144,7 @@ private Builder(QueryJobConfiguration jobConfiguration) { this.defaultDataset = jobConfiguration.defaultDataset; this.priority = jobConfiguration.priority; this.allowLargeResults = jobConfiguration.allowLargeResults; + this.createSession = jobConfiguration.createSession; this.useQueryCache = jobConfiguration.useQueryCache; this.flattenResults = jobConfiguration.flattenResults; this.dryRun = jobConfiguration.dryRun; @@ -185,6 +188,7 @@ private Builder(com.google.api.services.bigquery.model.JobConfiguration configur } } allowLargeResults = queryConfigurationPb.getAllowLargeResults(); + createSession = queryConfigurationPb.getCreateSession(); useQueryCache = queryConfigurationPb.getUseQueryCache(); flattenResults = queryConfigurationPb.getFlattenResults(); useLegacySql = queryConfigurationPb.getUseLegacySql(); @@ -460,6 +464,16 @@ public Builder setPriority(Priority priority) { return this; } + /** + * Sets whether to create a new session. If {@code true} a random session id will be generated + * by BigQuery. If false, runs query with an existing session_id passed in ConnectionProperty, + * otherwise runs query in non-session mode." + */ + public Builder setCreateSession(Boolean createSession) { + this.createSession = createSession; + return this; + } + /** * Sets whether the job is enabled to create arbitrarily large results. If {@code true} the * query is allowed to create large results at a slight cost in performance. If {@code true} @@ -656,6 +670,7 @@ private QueryJobConfiguration(Builder builder) { namedParameters = ImmutableMap.copyOf(builder.namedParameters); this.parameterMode = builder.parameterMode; this.allowLargeResults = builder.allowLargeResults; + this.createSession = builder.createSession; this.createDisposition = builder.createDisposition; this.defaultDataset = builder.defaultDataset; this.destinationTable = builder.destinationTable; @@ -693,6 +708,15 @@ public Boolean allowLargeResults() { return allowLargeResults; } + /** + * Returns whether to create a new session. + * + * @see Create Sessions + */ + public Boolean createSession() { + return createSession; + } + /** * Returns whether the job is allowed to create new tables. * @@ -897,6 +921,7 @@ ToStringHelper toStringHelper() { .add("destinationEncryptionConfiguration", destinationEncryptionConfiguration) .add("defaultDataset", defaultDataset) .add("allowLargeResults", allowLargeResults) + .add("createSession", createSession) .add("flattenResults", flattenResults) .add("priority", priority) .add("tableDefinitions", tableDefinitions) @@ -928,6 +953,7 @@ public int hashCode() { return Objects.hash( baseHashCode(), allowLargeResults, + createSession, createDisposition, destinationTable, defaultDataset, @@ -988,6 +1014,9 @@ com.google.api.services.bigquery.model.JobConfiguration toPb() { if (allowLargeResults != null) { queryConfigurationPb.setAllowLargeResults(allowLargeResults); } + if (createSession != null) { + queryConfigurationPb.setCreateSession(createSession); + } if (createDisposition != null) { queryConfigurationPb.setCreateDisposition(createDisposition.toString()); } diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java index 7492809d6..ccde465af 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java @@ -37,6 +37,7 @@ final class QueryRequestInfo { private final String query; private final List queryParameters; private final String requestId; + private final Boolean createSession; private final Boolean useQueryCache; private final Boolean useLegacySql; @@ -51,6 +52,7 @@ final class QueryRequestInfo { this.query = config.getQuery(); this.queryParameters = config.toPb().getQuery().getQueryParameters(); this.requestId = UUID.randomUUID().toString(); + this.createSession = config.createSession(); this.useLegacySql = config.useLegacySql(); this.useQueryCache = config.useQueryCache(); } @@ -97,6 +99,9 @@ QueryRequest toPb() { if (queryParameters != null) { request.setQueryParameters(queryParameters); } + if (createSession != null) { + request.setCreateSession(createSession); + } if (useLegacySql != null) { request.setUseLegacySql(useLegacySql); } @@ -118,6 +123,7 @@ public String toString() { .add("query", query) .add("requestId", requestId) .add("queryParameters", queryParameters) + .add("createSession", createSession) .add("useQueryCache", useQueryCache) .add("useLegacySql", useLegacySql) .toString(); @@ -135,6 +141,7 @@ public int hashCode() { query, queryParameters, requestId, + createSession, useQueryCache, useLegacySql); } diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobStatisticsTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobStatisticsTest.java index c421c0370..87fae41b5 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobStatisticsTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobStatisticsTest.java @@ -26,6 +26,7 @@ import com.google.cloud.bigquery.JobStatistics.ReservationUsage; import com.google.cloud.bigquery.JobStatistics.ScriptStatistics; import com.google.cloud.bigquery.JobStatistics.ScriptStatistics.ScriptStackFrame; +import com.google.cloud.bigquery.JobStatistics.SessionInfo; import com.google.cloud.bigquery.JobStatistics.TransactionInfo; import com.google.cloud.bigquery.QueryStage.QueryStep; import com.google.common.collect.ImmutableList; @@ -73,6 +74,7 @@ public class JobStatisticsTest { private static final String NAME = "reservation-name"; private static final Long SLOTMS = 12545L; private static final String TRANSACTION_ID = UUID.randomUUID().toString().substring(0, 8); + private static final String SESSION_ID = UUID.randomUUID().toString().substring(0, 8); private static final CopyStatistics COPY_STATISTICS = CopyStatistics.newBuilder() .setCreationTimestamp(CREATION_TIME) @@ -222,6 +224,9 @@ public class JobStatisticsTest { private static final TransactionInfo TRANSACTION_INFO = TransactionInfo.newbuilder().setTransactionId(TRANSACTION_ID).build(); + private static final SessionInfo SESSION_INFO = + SessionInfo.newBuilder().setSessionId(SESSION_ID).build(); + @Test public void testBuilder() { assertEquals(CREATION_TIME, EXTRACT_STATISTICS.getCreationTime()); @@ -293,6 +298,7 @@ public void testBuilder() { assertEquals(NAME, RESERVATION_USAGE.getName()); assertEquals(SLOTMS, RESERVATION_USAGE.getSlotMs()); assertEquals(TRANSACTION_ID, TRANSACTION_INFO.getTransactionId()); + assertEquals(SESSION_ID, SESSION_INFO.getSessionId()); } @Test @@ -319,6 +325,7 @@ public void testToPbAndFromPb() { } compareReservation(RESERVATION_USAGE, ReservationUsage.fromPb(RESERVATION_USAGE.toPb())); compareTransactionInfo(TRANSACTION_INFO, TransactionInfo.fromPb(TRANSACTION_INFO.toPb())); + compareSessionInfo(SESSION_INFO, SessionInfo.fromPb(SESSION_INFO.toPb())); } @Test @@ -441,4 +448,12 @@ private void compareTransactionInfo(TransactionInfo expected, TransactionInfo va assertEquals(expected.toPb(), value.toPb()); assertEquals(expected.getTransactionId(), value.getTransactionId()); } + + private void compareSessionInfo(SessionInfo expected, SessionInfo value) { + assertEquals(expected, value); + assertEquals(expected.hashCode(), value.hashCode()); + assertEquals(expected.toString(), value.toString()); + assertEquals(expected.toPb(), value.toPb()); + assertEquals(expected.getSessionId(), value.getSessionId()); + } } diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryJobConfigurationTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryJobConfigurationTest.java index bcfc72356..9a20219d6 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryJobConfigurationTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryJobConfigurationTest.java @@ -78,6 +78,7 @@ public class QueryJobConfigurationTest { private static final Priority PRIORITY = Priority.BATCH; private static final boolean ALLOW_LARGE_RESULTS = true; private static final boolean USE_QUERY_CACHE = false; + private static final boolean CREATE_SESSION = true; private static final boolean FLATTEN_RESULTS = true; private static final boolean USE_LEGACY_SQL = true; private static final Integer MAX_BILLING_TIER = 123; @@ -119,6 +120,7 @@ public class QueryJobConfigurationTest { .setDestinationTable(TABLE_ID) .setWriteDisposition(WRITE_DISPOSITION) .setPriority(PRIORITY) + .setCreateSession(CREATE_SESSION) .setFlattenResults(FLATTEN_RESULTS) .setUserDefinedFunctions(USER_DEFINED_FUNCTIONS) .setDryRun(true) @@ -238,6 +240,7 @@ private void compareQueryJobConfiguration( assertEquals(expected.getCreateDisposition(), value.getCreateDisposition()); assertEquals(expected.getDefaultDataset(), value.getDefaultDataset()); assertEquals(expected.getDestinationTable(), value.getDestinationTable()); + assertEquals(expected.createSession(), value.createSession()); assertEquals(expected.flattenResults(), value.flattenResults()); assertEquals(expected.getPriority(), value.getPriority()); assertEquals(expected.getQuery(), value.getQuery()); diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryRequestInfoTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryRequestInfoTest.java index 4adeeaf00..5c224af80 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryRequestInfoTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryRequestInfoTest.java @@ -74,6 +74,7 @@ public class QueryRequestInfoTest { private static final WriteDisposition WRITE_DISPOSITION = WriteDisposition.WRITE_APPEND; private static final Priority PRIORITY = Priority.BATCH; private static final boolean ALLOW_LARGE_RESULTS = true; + private static final boolean CREATE_SESSION = true; private static final boolean USE_QUERY_CACHE = false; private static final boolean FLATTEN_RESULTS = true; private static final boolean USE_LEGACY_SQL = true; @@ -142,6 +143,7 @@ public class QueryRequestInfoTest { .setLabels(LABELS) .setConnectionProperties(CONNECTION_PROPERTIES) .setPositionalParameters(POSITIONAL_PARAMETER) + .setCreateSession(CREATE_SESSION) .setMaxResults(100L) .build(); QueryRequestInfo REQUEST_INFO_SUPPORTED = new QueryRequestInfo(QUERY_JOB_CONFIGURATION_SUPPORTED); diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java index 8ef74458f..88695f873 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java @@ -2235,6 +2235,39 @@ public void testFastQueryHTTPException() throws InterruptedException { } } + @Test + public void testQuerySessionSupport() throws InterruptedException { + String query = "CREATE TEMPORARY TABLE temptable AS SELECT 17 as foo"; + QueryJobConfiguration queryJobConfiguration = + QueryJobConfiguration.newBuilder(query) + .setDefaultDataset(DatasetId.of(DATASET)) + .setCreateSession(true) + .build(); + Job remoteJob = bigquery.create(JobInfo.of(queryJobConfiguration)); + remoteJob = remoteJob.waitFor(); + assertNull(remoteJob.getStatus().getError()); + + Job queryJob = bigquery.getJob(remoteJob.getJobId()); + JobStatistics.QueryStatistics statistics = queryJob.getStatistics(); + String sessionId = statistics.getSessionInfo().getSessionId(); + assertNotNull(sessionId); + + String queryTempTable = "SELECT * FROM temptable"; + ConnectionProperty connectionProperty = + ConnectionProperty.newBuilder().setKey("session_id").setValue(sessionId).build(); + QueryJobConfiguration queryJobConfigurationWithSession = + QueryJobConfiguration.newBuilder(queryTempTable) + .setDefaultDataset(DatasetId.of(DATASET)) + .setConnectionProperties(ImmutableList.of(connectionProperty)) + .build(); + Job remoteJobWithSession = bigquery.create(JobInfo.of(queryJobConfigurationWithSession)); + remoteJobWithSession = remoteJobWithSession.waitFor(); + assertNull(remoteJobWithSession.getStatus().getError()); + Job queryJobWithSession = bigquery.getJob(remoteJobWithSession.getJobId()); + JobStatistics.QueryStatistics statisticsWithSession = queryJobWithSession.getStatistics(); + assertEquals(sessionId, statisticsWithSession.getSessionInfo().getSessionId()); + } + @Test public void testDmlStatistics() throws InterruptedException { String tableName = TABLE_ID_FASTQUERY.getTable();