Skip to content

Commit 42851d8

Browse files
Neenu1995gcf-owl-bot[bot]
Neenu1995
and
gcf-owl-bot[bot]
authored
feat: add support for session id on TableDataWriteChannel (#2715)
* feat: add support for session id on TableDataWriteChannel * chore: change csv file name * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * chore: remove struct tests * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * chore: fix graalvm tests * chore: make path absolute * chore: try streams * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * chore: add IOUtils * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * chore: nit * cleanup --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
1 parent 5c64797 commit 42851d8

File tree

5 files changed

+224
-4
lines changed

5 files changed

+224
-4
lines changed

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

Lines changed: 49 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import com.google.cloud.bigquery.JobInfo.WriteDisposition;
2626
import com.google.common.base.MoreObjects;
2727
import com.google.common.collect.ImmutableList;
28+
import com.google.common.collect.Lists;
2829
import com.google.common.primitives.Ints;
2930
import java.io.Serializable;
3031
import java.util.List;
@@ -56,9 +57,11 @@ public final class WriteChannelConfiguration implements LoadConfiguration, Seria
5657
private final Boolean useAvroLogicalTypes;
5758
private final Map<String, String> labels;
5859
private List<String> decimalTargetTypes;
60+
private final List<ConnectionProperty> connectionProperties;
5961

60-
public static final class Builder implements LoadConfiguration.Builder {
62+
private final Boolean createSession;
6163

64+
public static final class Builder implements LoadConfiguration.Builder {
6265
private TableId destinationTable;
6366
private CreateDisposition createDisposition;
6467
private WriteDisposition writeDisposition;
@@ -75,6 +78,9 @@ public static final class Builder implements LoadConfiguration.Builder {
7578
private Boolean useAvroLogicalTypes;
7679
private Map<String, String> labels;
7780
private List<String> decimalTargetTypes;
81+
private List<ConnectionProperty> connectionProperties;
82+
83+
private Boolean createSession;
7884

7985
private Builder() {}
8086

@@ -96,6 +102,8 @@ private Builder(WriteChannelConfiguration writeChannelConfiguration) {
96102
this.useAvroLogicalTypes = writeChannelConfiguration.useAvroLogicalTypes;
97103
this.labels = writeChannelConfiguration.labels;
98104
this.decimalTargetTypes = writeChannelConfiguration.decimalTargetTypes;
105+
this.connectionProperties = writeChannelConfiguration.connectionProperties;
106+
this.createSession = writeChannelConfiguration.createSession;
99107
}
100108

101109
private Builder(com.google.api.services.bigquery.model.JobConfiguration configurationPb) {
@@ -175,6 +183,13 @@ private Builder(com.google.api.services.bigquery.model.JobConfiguration configur
175183
if (loadConfigurationPb.getDecimalTargetTypes() != null) {
176184
this.decimalTargetTypes = loadConfigurationPb.getDecimalTargetTypes();
177185
}
186+
if (loadConfigurationPb.getConnectionProperties() != null) {
187+
188+
this.connectionProperties =
189+
Lists.transform(
190+
loadConfigurationPb.getConnectionProperties(), ConnectionProperty.FROM_PB_FUNCTION);
191+
}
192+
createSession = loadConfigurationPb.getCreateSession();
178193
}
179194

180195
@Override
@@ -274,6 +289,16 @@ public Builder setDecimalTargetTypes(List<String> decimalTargetTypes) {
274289
return this;
275290
}
276291

292+
public Builder setConnectionProperties(List<ConnectionProperty> connectionProperties) {
293+
this.connectionProperties = ImmutableList.copyOf(connectionProperties);
294+
return this;
295+
}
296+
297+
public Builder setCreateSession(Boolean createSession) {
298+
this.createSession = createSession;
299+
return this;
300+
}
301+
277302
@Override
278303
public WriteChannelConfiguration build() {
279304
return new WriteChannelConfiguration(this);
@@ -297,6 +322,8 @@ protected WriteChannelConfiguration(Builder builder) {
297322
this.useAvroLogicalTypes = builder.useAvroLogicalTypes;
298323
this.labels = builder.labels;
299324
this.decimalTargetTypes = builder.decimalTargetTypes;
325+
this.connectionProperties = builder.connectionProperties;
326+
this.createSession = builder.createSession;
300327
}
301328

302329
@Override
@@ -390,6 +417,14 @@ public List<String> getDecimalTargetTypes() {
390417
return decimalTargetTypes;
391418
}
392419

420+
public List<ConnectionProperty> getConnectionProperties() {
421+
return connectionProperties;
422+
}
423+
424+
public Boolean getCreateSession() {
425+
return createSession;
426+
}
427+
393428
@Override
394429
public Builder toBuilder() {
395430
return new Builder(this);
@@ -412,7 +447,9 @@ MoreObjects.ToStringHelper toStringHelper() {
412447
.add("clustering", clustering)
413448
.add("useAvroLogicalTypes", useAvroLogicalTypes)
414449
.add("labels", labels)
415-
.add("decimalTargetTypes", decimalTargetTypes);
450+
.add("decimalTargetTypes", decimalTargetTypes)
451+
.add("connectionProperties", connectionProperties)
452+
.add("createSession", createSession);
416453
}
417454

418455
@Override
@@ -444,7 +481,9 @@ public int hashCode() {
444481
clustering,
445482
useAvroLogicalTypes,
446483
labels,
447-
decimalTargetTypes);
484+
decimalTargetTypes,
485+
connectionProperties,
486+
createSession);
448487
}
449488

450489
WriteChannelConfiguration setProjectId(String projectId) {
@@ -519,6 +558,13 @@ com.google.api.services.bigquery.model.JobConfiguration toPb() {
519558
if (decimalTargetTypes != null) {
520559
loadConfigurationPb.setDecimalTargetTypes(decimalTargetTypes);
521560
}
561+
if (connectionProperties != null) {
562+
loadConfigurationPb.setConnectionProperties(
563+
Lists.transform(connectionProperties, ConnectionProperty.TO_PB_FUNCTION));
564+
}
565+
if (createSession != null) {
566+
loadConfigurationPb.setCreateSession(createSession);
567+
}
522568
jobConfiguration.setLoad(loadConfigurationPb);
523569
return jobConfiguration;
524570
}

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,14 @@ public class WriteChannelConfigurationTest {
6161
ImmutableMap.of("test-job-name", "test-write-channel");
6262
private static final List<String> DECIMAL_TARGET_TYPES =
6363
ImmutableList.of("NUMERIC", "BIGNUMERIC");
64+
65+
private static final boolean CREATE_SESSION = true;
66+
private static final String KEY = "session_id";
67+
private static final String VALUE = "session_id_1234567890";
68+
private static final ConnectionProperty CONNECTION_PROPERTY =
69+
ConnectionProperty.newBuilder().setKey(KEY).setValue(VALUE).build();
70+
private static final List<ConnectionProperty> CONNECTION_PROPERTIES =
71+
ImmutableList.of(CONNECTION_PROPERTY);
6472
private static final WriteChannelConfiguration LOAD_CONFIGURATION_CSV =
6573
WriteChannelConfiguration.newBuilder(TABLE_ID)
6674
.setCreateDisposition(CREATE_DISPOSITION)
@@ -76,6 +84,8 @@ public class WriteChannelConfigurationTest {
7684
.setClustering(CLUSTERING)
7785
.setLabels(LABELS)
7886
.setDecimalTargetTypes(DECIMAL_TARGET_TYPES)
87+
.setConnectionProperties(CONNECTION_PROPERTIES)
88+
.setCreateSession(CREATE_SESSION)
7989
.build();
8090

8191
private static final DatastoreBackupOptions BACKUP_OPTIONS =
@@ -232,5 +242,7 @@ private void compareLoadConfiguration(
232242
assertEquals(expected.getUseAvroLogicalTypes(), value.getUseAvroLogicalTypes());
233243
assertEquals(expected.getLabels(), value.getLabels());
234244
assertEquals(expected.getDecimalTargetTypes(), value.getDecimalTargetTypes());
245+
assertEquals(expected.getConnectionProperties(), value.getConnectionProperties());
246+
assertEquals(expected.getCreateSession(), value.getCreateSession());
235247
}
236248
}

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

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import static org.junit.Assert.assertTrue;
2929
import static org.junit.Assert.fail;
3030

31+
import com.google.api.client.util.IOUtils;
3132
import com.google.api.gax.paging.Page;
3233
import com.google.auth.oauth2.GoogleCredentials;
3334
import com.google.auth.oauth2.ServiceAccountCredentials;
@@ -60,6 +61,7 @@
6061
import com.google.cloud.bigquery.ConnectionProperty;
6162
import com.google.cloud.bigquery.ConnectionSettings;
6263
import com.google.cloud.bigquery.CopyJobConfiguration;
64+
import com.google.cloud.bigquery.CsvOptions;
6365
import com.google.cloud.bigquery.Dataset;
6466
import com.google.cloud.bigquery.DatasetId;
6567
import com.google.cloud.bigquery.DatasetInfo;
@@ -141,9 +143,13 @@
141143
import com.google.gson.JsonObject;
142144
import java.io.IOException;
143145
import java.io.InputStream;
146+
import java.io.OutputStream;
144147
import java.math.BigDecimal;
145148
import java.nio.ByteBuffer;
149+
import java.nio.channels.Channels;
146150
import java.nio.charset.StandardCharsets;
151+
import java.nio.file.FileSystems;
152+
import java.nio.file.Path;
147153
import java.sql.ResultSet;
148154
import java.sql.SQLException;
149155
import java.sql.Time;
@@ -711,6 +717,36 @@ public class ITBigQueryTest {
711717
private static final List<ConnectionProperty> CONNECTION_PROPERTIES =
712718
ImmutableList.of(CONNECTION_PROPERTY);
713719

720+
private static final Field ID_SCHEMA =
721+
Field.newBuilder("id", LegacySQLTypeName.STRING)
722+
.setMode(Mode.REQUIRED)
723+
.setDescription("id")
724+
.build();
725+
private static final Field FIRST_NAME_SCHEMA =
726+
Field.newBuilder("firstname", LegacySQLTypeName.STRING)
727+
.setMode(Field.Mode.NULLABLE)
728+
.setDescription("First Name")
729+
.build();
730+
private static final Field LAST_NAME_SCHEMA =
731+
Field.newBuilder("lastname", LegacySQLTypeName.STRING)
732+
.setMode(Field.Mode.NULLABLE)
733+
.setDescription("LAST NAME")
734+
.build();
735+
private static final Field EMAIL_SCHEMA =
736+
Field.newBuilder("email", LegacySQLTypeName.STRING)
737+
.setMode(Field.Mode.NULLABLE)
738+
.setDescription("email")
739+
.build();
740+
private static final Field PROFESSION_SCHEMA =
741+
Field.newBuilder("profession", LegacySQLTypeName.STRING)
742+
.setMode(Field.Mode.NULLABLE)
743+
.setDescription("profession")
744+
.build();
745+
private static final Schema SESSION_TABLE_SCHEMA =
746+
Schema.of(ID_SCHEMA, FIRST_NAME_SCHEMA, LAST_NAME_SCHEMA, EMAIL_SCHEMA, PROFESSION_SCHEMA);
747+
private static final Path csvPath =
748+
FileSystems.getDefault().getPath("src/test/resources", "sessionTest.csv").toAbsolutePath();
749+
714750
private static final Set<String> PUBLIC_DATASETS =
715751
ImmutableSet.of("github_repos", "hacker_news", "noaa_gsod", "samples", "usa_names");
716752

@@ -3733,6 +3769,80 @@ public void testQuerySessionSupport() throws InterruptedException {
37333769
assertEquals(sessionId, statisticsWithSession.getSessionInfo().getSessionId());
37343770
}
37353771

3772+
@Test
3773+
public void testLoadSessionSupportWriteChannelConfiguration() throws InterruptedException {
3774+
TableId sessionTableId = TableId.of("_SESSION", "test_temp_destination_table_from_file");
3775+
3776+
WriteChannelConfiguration configuration =
3777+
WriteChannelConfiguration.newBuilder(sessionTableId)
3778+
.setFormatOptions(CsvOptions.newBuilder().setFieldDelimiter(",").build())
3779+
.setCreateDisposition(JobInfo.CreateDisposition.CREATE_IF_NEEDED)
3780+
.setSchema(SESSION_TABLE_SCHEMA)
3781+
.setCreateSession(true)
3782+
.build();
3783+
String jobName = "jobId_" + UUID.randomUUID().toString();
3784+
JobId jobId = JobId.newBuilder().setLocation("us").setJob(jobName).build();
3785+
String sessionId;
3786+
3787+
// Imports a local file into a table.
3788+
try (TableDataWriteChannel writer = bigquery.writer(jobId, configuration);
3789+
OutputStream stream = Channels.newOutputStream(writer)) {
3790+
InputStream inputStream =
3791+
ITBigQueryTest.class.getClassLoader().getResourceAsStream("sessionTest.csv");
3792+
// Can use `Files.copy(csvPath, stream);` instead.
3793+
// Using IOUtils here because graalvm can't handle resource files.
3794+
IOUtils.copy(inputStream, stream);
3795+
3796+
} catch (IOException e) {
3797+
throw new RuntimeException(e);
3798+
}
3799+
Job loadJob = bigquery.getJob(jobId);
3800+
Job completedJob = loadJob.waitFor();
3801+
3802+
assertNotNull(completedJob);
3803+
assertEquals(jobId.getJob(), completedJob.getJobId().getJob());
3804+
JobStatistics.LoadStatistics statistics = completedJob.getStatistics();
3805+
3806+
sessionId = statistics.getSessionInfo().getSessionId();
3807+
assertNotNull(sessionId);
3808+
3809+
// Load job in the same session.
3810+
// Should load the data to a temp table.
3811+
ConnectionProperty sessionConnectionProperty =
3812+
ConnectionProperty.newBuilder().setKey("session_id").setValue(sessionId).build();
3813+
WriteChannelConfiguration sessionConfiguration =
3814+
WriteChannelConfiguration.newBuilder(sessionTableId)
3815+
.setConnectionProperties(ImmutableList.of(sessionConnectionProperty))
3816+
.setFormatOptions(CsvOptions.newBuilder().setFieldDelimiter(",").build())
3817+
.setCreateDisposition(JobInfo.CreateDisposition.CREATE_IF_NEEDED)
3818+
.setSchema(SESSION_TABLE_SCHEMA)
3819+
.build();
3820+
String sessionJobName = "jobId_" + UUID.randomUUID().toString();
3821+
JobId sessionJobId = JobId.newBuilder().setLocation("us").setJob(sessionJobName).build();
3822+
try (TableDataWriteChannel writer = bigquery.writer(sessionJobId, sessionConfiguration);
3823+
OutputStream stream = Channels.newOutputStream(writer)) {
3824+
InputStream inputStream =
3825+
ITBigQueryTest.class.getClassLoader().getResourceAsStream("sessionTest.csv");
3826+
IOUtils.copy(inputStream, stream);
3827+
} catch (IOException e) {
3828+
throw new RuntimeException(e);
3829+
}
3830+
Job queryJobWithSession = bigquery.getJob(sessionJobId);
3831+
queryJobWithSession = queryJobWithSession.waitFor();
3832+
LoadStatistics statisticsWithSession = queryJobWithSession.getStatistics();
3833+
assertNotNull(statisticsWithSession.getSessionInfo().getSessionId());
3834+
3835+
// Checking if the data loaded to the temp table in the session
3836+
String queryTempTable = "SELECT * FROM _SESSION.test_temp_destination_table_from_file;";
3837+
QueryJobConfiguration queryJobConfigurationWithSession =
3838+
QueryJobConfiguration.newBuilder(queryTempTable)
3839+
.setConnectionProperties(ImmutableList.of(sessionConnectionProperty))
3840+
.build();
3841+
Job queryTempTableJob = bigquery.create(JobInfo.of(queryJobConfigurationWithSession));
3842+
queryTempTableJob = queryTempTableJob.waitFor();
3843+
assertNotNull(queryTempTableJob.getQueryResults());
3844+
}
3845+
37363846
@Test
37373847
public void testLoadSessionSupport() throws InterruptedException {
37383848
// Start the session

‎google-cloud-bigquery/src/test/resources/META-INF/native-image/resource-config.json

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
11
{
2-
"resources":[{"pattern": ".*.csv"}]
2+
"resources":[{"pattern": ".*.csv"},
3+
{"pattern": ".*src/test/resources/sessionTest.csv"}]
34
}

‎google-cloud-bigquery/src/test/resources/sessionTest.csv

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
id,firstname,lastname,email,profession
2+
100,Rani,Merell,[email protected],firefighter
3+
101,Goldie,Dex,[email protected],developer
4+
102,Cristabel,Munn,[email protected],developer
5+
103,Genevra,Strephon,[email protected],firefighter
6+
104,Augustine,Thema,[email protected],doctor
7+
105,Jemie,Gombach,[email protected],police officer
8+
106,Maye,Stuart,[email protected],developer
9+
107,Ayn,Carmena,[email protected],worker
10+
108,Gale,Celestine,[email protected],doctor
11+
109,Alex,Jerold,[email protected],firefighter
12+
110,Violet,Giule,[email protected],firefighter
13+
111,Starla,Uird,[email protected],doctor
14+
112,Tarra,Pelagias,[email protected],police officer
15+
113,Eugine,Deny,[email protected],doctor
16+
114,Shirlee,Ricarda,[email protected],doctor
17+
115,Ariela,Penelopa,[email protected],worker
18+
116,Lelah,Astra,[email protected],police officer
19+
117,Debee,Deegan,[email protected],developer
20+
118,Pollyanna,Euridice,[email protected],worker
21+
119,Cathie,Halsey,[email protected],firefighter
22+
120,Rebeca,Quinn,[email protected],doctor
23+
121,Paulita,Arquit,[email protected],police officer
24+
122,Rebeca,Emanuel,[email protected],firefighter
25+
123,Tera,Ilka,[email protected],firefighter
26+
124,Orsola,Briney,[email protected],doctor
27+
125,Paulita,Wyn,[email protected],doctor
28+
126,Constance,Christine,[email protected],firefighter
29+
127,Claresta,Kinnard,[email protected],developer
30+
128,Leanna,Mendez,[email protected],developer
31+
129,Corina,Chabot,[email protected],developer
32+
130,Romona,Audly,[email protected],worker
33+
131,Cordi,Lynn,[email protected],firefighter
34+
132,Sheree,Tyson,[email protected],worker
35+
133,Jinny,Bevin,[email protected],police officer
36+
134,Kassey,Havens,[email protected],firefighter
37+
135,Wanda,Thema,[email protected],developer
38+
136,Vita,Jagir,[email protected],developer
39+
137,Alie,Aprile,[email protected],firefighter
40+
138,Modestia,Jena,[email protected],doctor
41+
139,Cyndie,Pelagias,[email protected],worker
42+
140,Ariela,Lilybelle,[email protected],firefighter
43+
141,Jan,Parette,[email protected],firefighter
44+
142,Merry,Horan,[email protected],developer
45+
143,Katuscha,Candy,[email protected],police officer
46+
144,Kerrin,Heisel,[email protected],developer
47+
145,Nollie,Magdalen,[email protected],doctor
48+
146,Karlee,Gordon,[email protected],developer
49+
147,Dolli,Fadiman,[email protected],firefighter
50+
148,Leontine,Delp,[email protected],worker
51+
149,Ricky,Nadia,[email protected],doctor

0 commit comments

Comments
 (0)