Skip to content
This repository was archived by the owner on Mar 23, 2026. It is now read-only.

Commit e431c17

Browse files
Neenu1995gcf-owl-bot[bot]
Neenu1995
and
gcf-owl-bot[bot]
authored
feat: add support for session_id in load jobs (#2519)
* feat: add support for session_id in load jobs * chore: fix variable change * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
1 parent 59b933e commit e431c17

3 files changed

Lines changed: 109 additions & 1 deletion

File tree

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

Lines changed: 46 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import com.google.common.base.MoreObjects.ToStringHelper;
2323
import com.google.common.base.Strings;
2424
import com.google.common.collect.ImmutableList;
25+
import com.google.common.collect.Lists;
2526
import com.google.common.primitives.Ints;
2627
import java.util.List;
2728
import java.util.Map;
@@ -58,6 +59,10 @@ public final class LoadJobConfiguration extends JobConfiguration implements Load
5859
private final HivePartitioningOptions hivePartitioningOptions;
5960
private final String referenceFileSchemaUri;
6061

62+
private final List<ConnectionProperty> connectionProperties;
63+
64+
private final Boolean createSession;
65+
6166
public static final class Builder extends JobConfiguration.Builder<LoadJobConfiguration, Builder>
6267
implements LoadConfiguration.Builder {
6368

@@ -83,6 +88,8 @@ public static final class Builder extends JobConfiguration.Builder<LoadJobConfig
8388
private RangePartitioning rangePartitioning;
8489
private HivePartitioningOptions hivePartitioningOptions;
8590
private String referenceFileSchemaUri;
91+
private List<ConnectionProperty> connectionProperties;
92+
private Boolean createSession;
8693

8794
private Builder() {
8895
super(Type.LOAD);
@@ -112,6 +119,8 @@ private Builder(LoadJobConfiguration loadConfiguration) {
112119
this.rangePartitioning = loadConfiguration.rangePartitioning;
113120
this.hivePartitioningOptions = loadConfiguration.hivePartitioningOptions;
114121
this.referenceFileSchemaUri = loadConfiguration.referenceFileSchemaUri;
122+
this.connectionProperties = loadConfiguration.connectionProperties;
123+
this.createSession = loadConfiguration.createSession;
115124
}
116125

117126
private Builder(com.google.api.services.bigquery.model.JobConfiguration configurationPb) {
@@ -205,6 +214,13 @@ private Builder(com.google.api.services.bigquery.model.JobConfiguration configur
205214
if (loadConfigurationPb.getReferenceFileSchemaUri() != null) {
206215
this.referenceFileSchemaUri = loadConfigurationPb.getReferenceFileSchemaUri();
207216
}
217+
if (loadConfigurationPb.getConnectionProperties() != null) {
218+
219+
this.connectionProperties =
220+
Lists.transform(
221+
loadConfigurationPb.getConnectionProperties(), ConnectionProperty.FROM_PB_FUNCTION);
222+
}
223+
createSession = loadConfigurationPb.getCreateSession();
208224
}
209225

210226
@Override
@@ -368,6 +384,16 @@ public Builder setReferenceFileSchemaUri(String referenceFileSchemaUri) {
368384
return this;
369385
}
370386

387+
public Builder setConnectionProperties(List<ConnectionProperty> connectionProperties) {
388+
this.connectionProperties = ImmutableList.copyOf(connectionProperties);
389+
return this;
390+
}
391+
392+
public Builder setCreateSession(Boolean createSession) {
393+
this.createSession = createSession;
394+
return this;
395+
}
396+
371397
@Override
372398
public LoadJobConfiguration build() {
373399
return new LoadJobConfiguration(this);
@@ -397,6 +423,8 @@ private LoadJobConfiguration(Builder builder) {
397423
this.rangePartitioning = builder.rangePartitioning;
398424
this.hivePartitioningOptions = builder.hivePartitioningOptions;
399425
this.referenceFileSchemaUri = builder.referenceFileSchemaUri;
426+
this.connectionProperties = builder.connectionProperties;
427+
this.createSession = builder.createSession;
400428
}
401429

402430
@Override
@@ -520,6 +548,14 @@ public String getReferenceFileSchemaUri() {
520548
return referenceFileSchemaUri;
521549
}
522550

551+
public List<ConnectionProperty> getConnectionProperties() {
552+
return connectionProperties;
553+
}
554+
555+
public Boolean getCreateSession() {
556+
return createSession;
557+
}
558+
523559
@Override
524560
public Builder toBuilder() {
525561
return new Builder(this);
@@ -548,7 +584,9 @@ ToStringHelper toStringHelper() {
548584
.add("jobTimeoutMs", jobTimeoutMs)
549585
.add("rangePartitioning", rangePartitioning)
550586
.add("hivePartitioningOptions", hivePartitioningOptions)
551-
.add("referenceFileSchemaUri", referenceFileSchemaUri);
587+
.add("referenceFileSchemaUri", referenceFileSchemaUri)
588+
.add("connectionProperties", connectionProperties)
589+
.add("createSession", createSession);
552590
}
553591

554592
@Override
@@ -654,6 +692,13 @@ com.google.api.services.bigquery.model.JobConfiguration toPb() {
654692
if (referenceFileSchemaUri != null) {
655693
loadConfigurationPb.setReferenceFileSchemaUri(referenceFileSchemaUri);
656694
}
695+
if (connectionProperties != null) {
696+
loadConfigurationPb.setConnectionProperties(
697+
Lists.transform(connectionProperties, ConnectionProperty.TO_PB_FUNCTION));
698+
}
699+
if (createSession != null) {
700+
loadConfigurationPb.setCreateSession(createSession);
701+
}
657702

658703
jobConfiguration.setLoad(loadConfigurationPb);
659704
return jobConfiguration;

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,8 @@ public class LoadJobConfigurationTest {
5757
private static final Schema TABLE_SCHEMA = Schema.of(FIELD_SCHEMA);
5858
private static final Boolean AUTODETECT = true;
5959
private static final Boolean USE_AVRO_LOGICAL_TYPES = true;
60+
61+
private static final boolean CREATE_SESSION = true;
6062
private static final EncryptionConfiguration JOB_ENCRYPTION_CONFIGURATION =
6163
EncryptionConfiguration.newBuilder().setKmsKeyName("KMS_KEY_1").build();
6264
private static final TimePartitioning TIME_PARTITIONING = TimePartitioning.of(Type.DAY);
@@ -71,6 +73,13 @@ public class LoadJobConfigurationTest {
7173
RangePartitioning.newBuilder().setField("IntegerField").setRange(RANGE).build();
7274
private static final String MODE = "STRING";
7375
private static final String SOURCE_URI_PREFIX = "gs://bucket/path_to_table";
76+
77+
private static final String KEY = "session_id";
78+
private static final String VALUE = "session_id_1234567890";
79+
private static final ConnectionProperty CONNECTION_PROPERTY =
80+
ConnectionProperty.newBuilder().setKey(KEY).setValue(VALUE).build();
81+
private static final List<ConnectionProperty> CONNECTION_PROPERTIES =
82+
ImmutableList.of(CONNECTION_PROPERTY);
7483
private static final HivePartitioningOptions HIVE_PARTITIONING_OPTIONS =
7584
HivePartitioningOptions.newBuilder()
7685
.setMode(MODE)
@@ -95,6 +104,8 @@ public class LoadJobConfigurationTest {
95104
.setRangePartitioning(RANGE_PARTITIONING)
96105
.setNullMarker("nullMarker")
97106
.setHivePartitioningOptions(HIVE_PARTITIONING_OPTIONS)
107+
.setConnectionProperties(CONNECTION_PROPERTIES)
108+
.setCreateSession(CREATE_SESSION)
98109
.build();
99110

100111
private static final DatastoreBackupOptions BACKUP_OPTIONS =
@@ -253,5 +264,7 @@ private void compareLoadJobConfiguration(
253264
assertEquals(expected.getRangePartitioning(), value.getRangePartitioning());
254265
assertEquals(expected.getNullMarker(), value.getNullMarker());
255266
assertEquals(expected.getHivePartitioningOptions(), value.getHivePartitioningOptions());
267+
assertEquals(expected.getConnectionProperties(), value.getConnectionProperties());
268+
assertEquals(expected.getCreateSession(), value.getCreateSession());
256269
}
257270
}

‎google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3655,6 +3655,56 @@ public void testQuerySessionSupport() throws InterruptedException {
36553655
assertEquals(sessionId, statisticsWithSession.getSessionInfo().getSessionId());
36563656
}
36573657

3658+
@Test
3659+
public void testLoadSessionSupport() throws InterruptedException {
3660+
// Start the session
3661+
TableId sessionTableId = TableId.of("_SESSION", "test_temp_destination_table");
3662+
LoadJobConfiguration configuration =
3663+
LoadJobConfiguration.newBuilder(
3664+
sessionTableId, "gs://" + BUCKET + "/" + JSON_LOAD_FILE, FormatOptions.json())
3665+
.setCreateDisposition(JobInfo.CreateDisposition.CREATE_IF_NEEDED)
3666+
.setSchema(TABLE_SCHEMA)
3667+
.setCreateSession(true)
3668+
.build();
3669+
Job job = bigquery.create(JobInfo.of(configuration));
3670+
job = job.waitFor();
3671+
assertNull(job.getStatus().getError());
3672+
3673+
Job loadJob = bigquery.getJob(job.getJobId());
3674+
JobStatistics.LoadStatistics statistics = loadJob.getStatistics();
3675+
String sessionId = statistics.getSessionInfo().getSessionId();
3676+
assertNotNull(sessionId);
3677+
3678+
// Load job in the same session.
3679+
// Should load the data to a temp table.
3680+
ConnectionProperty sessionConnectionProperty =
3681+
ConnectionProperty.newBuilder().setKey("session_id").setValue(sessionId).build();
3682+
LoadJobConfiguration loadJobConfigurationWithSession =
3683+
LoadJobConfiguration.newBuilder(
3684+
sessionTableId, "gs://" + BUCKET + "/" + JSON_LOAD_FILE, FormatOptions.json())
3685+
.setCreateDisposition(JobInfo.CreateDisposition.CREATE_IF_NEEDED)
3686+
.setSchema(TABLE_SCHEMA)
3687+
.setConnectionProperties(ImmutableList.of(sessionConnectionProperty))
3688+
.build();
3689+
Job remoteJobWithSession = bigquery.create(JobInfo.of(loadJobConfigurationWithSession));
3690+
remoteJobWithSession = remoteJobWithSession.waitFor();
3691+
assertNull(remoteJobWithSession.getStatus().getError());
3692+
Job queryJobWithSession = bigquery.getJob(remoteJobWithSession.getJobId());
3693+
LoadStatistics statisticsWithSession = queryJobWithSession.getStatistics();
3694+
assertNotNull(statisticsWithSession.getSessionInfo().getSessionId());
3695+
3696+
// Checking if the data loaded to the temp table in the session
3697+
String queryTempTable = "SELECT * FROM _SESSION.test_temp_destination_table;";
3698+
QueryJobConfiguration queryJobConfigurationWithSession =
3699+
QueryJobConfiguration.newBuilder(queryTempTable)
3700+
.setConnectionProperties(ImmutableList.of(sessionConnectionProperty))
3701+
.build();
3702+
Job queryTempTableJob = bigquery.create(JobInfo.of(queryJobConfigurationWithSession));
3703+
queryTempTableJob = queryTempTableJob.waitFor();
3704+
assertNull(queryTempTableJob.getStatus().getError());
3705+
assertNotNull(queryTempTableJob.getQueryResults());
3706+
}
3707+
36583708
// TODO: uncomment this testcase when executeUpdate is implemented
36593709
// @Test
36603710
// public void testExecuteSelectWithSession() throws BigQuerySQLException {

0 commit comments

Comments
 (0)