Skip to content

Commit

Permalink
feat: add support for session id on TableDataWriteChannel (#2715)
Browse files Browse the repository at this point in the history 
* 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>
  • Loading branch information
@Neenu1995 @gcf-owl-bot
Neenu1995 and gcf-owl-bot[bot] committed May 30, 2023
1 parent 5c64797 commit 42851d8
Show file tree
Hide file tree
Showing 5 changed files with 224 additions and 4 deletions.
52 changes: 49 additions & 3 deletions google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/WriteChannelConfiguration.java
View file
Expand Up @@ -25,6 +25,7 @@
import com.google.cloud.bigquery.JobInfo.WriteDisposition;
import com.google.common.base.MoreObjects;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Lists;
import com.google.common.primitives.Ints;
import java.io.Serializable;
import java.util.List;
Expand Down Expand Up @@ -56,9 +57,11 @@ public final class WriteChannelConfiguration implements LoadConfiguration, Seria
private final Boolean useAvroLogicalTypes;
private final Map<String, String> labels;
private List<String> decimalTargetTypes;
private final List<ConnectionProperty> connectionProperties;

public static final class Builder implements LoadConfiguration.Builder {
private final Boolean createSession;

public static final class Builder implements LoadConfiguration.Builder {
private TableId destinationTable;
private CreateDisposition createDisposition;
private WriteDisposition writeDisposition;
Expand All @@ -75,6 +78,9 @@ public static final class Builder implements LoadConfiguration.Builder {
private Boolean useAvroLogicalTypes;
private Map<String, String> labels;
private List<String> decimalTargetTypes;
private List<ConnectionProperty> connectionProperties;

private Boolean createSession;

private Builder() {}

Expand All @@ -96,6 +102,8 @@ private Builder(WriteChannelConfiguration writeChannelConfiguration) {
this.useAvroLogicalTypes = writeChannelConfiguration.useAvroLogicalTypes;
this.labels = writeChannelConfiguration.labels;
this.decimalTargetTypes = writeChannelConfiguration.decimalTargetTypes;
this.connectionProperties = writeChannelConfiguration.connectionProperties;
this.createSession = writeChannelConfiguration.createSession;
}

private Builder(com.google.api.services.bigquery.model.JobConfiguration configurationPb) {
Expand Down Expand Up @@ -175,6 +183,13 @@ private Builder(com.google.api.services.bigquery.model.JobConfiguration configur
if (loadConfigurationPb.getDecimalTargetTypes() != null) {
this.decimalTargetTypes = loadConfigurationPb.getDecimalTargetTypes();
}
if (loadConfigurationPb.getConnectionProperties() != null) {

this.connectionProperties =
Lists.transform(
loadConfigurationPb.getConnectionProperties(), ConnectionProperty.FROM_PB_FUNCTION);
}
createSession = loadConfigurationPb.getCreateSession();
}

@Override
Expand Down Expand Up @@ -274,6 +289,16 @@ public Builder setDecimalTargetTypes(List<String> decimalTargetTypes) {
return this;
}

public Builder setConnectionProperties(List<ConnectionProperty> connectionProperties) {
this.connectionProperties = ImmutableList.copyOf(connectionProperties);
return this;
}

public Builder setCreateSession(Boolean createSession) {
this.createSession = createSession;
return this;
}

@Override
public WriteChannelConfiguration build() {
return new WriteChannelConfiguration(this);
Expand All @@ -297,6 +322,8 @@ protected WriteChannelConfiguration(Builder builder) {
this.useAvroLogicalTypes = builder.useAvroLogicalTypes;
this.labels = builder.labels;
this.decimalTargetTypes = builder.decimalTargetTypes;
this.connectionProperties = builder.connectionProperties;
this.createSession = builder.createSession;
}

@Override
Expand Down Expand Up @@ -390,6 +417,14 @@ public List<String> getDecimalTargetTypes() {
return decimalTargetTypes;
}

public List<ConnectionProperty> getConnectionProperties() {
return connectionProperties;
}

public Boolean getCreateSession() {
return createSession;
}

@Override
public Builder toBuilder() {
return new Builder(this);
Expand All @@ -412,7 +447,9 @@ MoreObjects.ToStringHelper toStringHelper() {
.add("clustering", clustering)
.add("useAvroLogicalTypes", useAvroLogicalTypes)
.add("labels", labels)
.add("decimalTargetTypes", decimalTargetTypes);
.add("decimalTargetTypes", decimalTargetTypes)
.add("connectionProperties", connectionProperties)
.add("createSession", createSession);
}

@Override
Expand Down Expand Up @@ -444,7 +481,9 @@ public int hashCode() {
clustering,
useAvroLogicalTypes,
labels,
decimalTargetTypes);
decimalTargetTypes,
connectionProperties,
createSession);
}

WriteChannelConfiguration setProjectId(String projectId) {
Expand Down Expand Up @@ -519,6 +558,13 @@ com.google.api.services.bigquery.model.JobConfiguration toPb() {
if (decimalTargetTypes != null) {
loadConfigurationPb.setDecimalTargetTypes(decimalTargetTypes);
}
if (connectionProperties != null) {
loadConfigurationPb.setConnectionProperties(
Lists.transform(connectionProperties, ConnectionProperty.TO_PB_FUNCTION));
}
if (createSession != null) {
loadConfigurationPb.setCreateSession(createSession);
}
jobConfiguration.setLoad(loadConfigurationPb);
return jobConfiguration;
}
Expand Down
12 changes: 12 additions & 0 deletions ...cloud-bigquery/src/test/java/com/google/cloud/bigquery/WriteChannelConfigurationTest.java
View file
Expand Up @@ -61,6 +61,14 @@ public class WriteChannelConfigurationTest {
ImmutableMap.of("test-job-name", "test-write-channel");
private static final List<String> DECIMAL_TARGET_TYPES =
ImmutableList.of("NUMERIC", "BIGNUMERIC");

private static final boolean CREATE_SESSION = true;
private static final String KEY = "session_id";
private static final String VALUE = "session_id_1234567890";
private static final ConnectionProperty CONNECTION_PROPERTY =
ConnectionProperty.newBuilder().setKey(KEY).setValue(VALUE).build();
private static final List<ConnectionProperty> CONNECTION_PROPERTIES =
ImmutableList.of(CONNECTION_PROPERTY);
private static final WriteChannelConfiguration LOAD_CONFIGURATION_CSV =
WriteChannelConfiguration.newBuilder(TABLE_ID)
.setCreateDisposition(CREATE_DISPOSITION)
Expand All @@ -76,6 +84,8 @@ public class WriteChannelConfigurationTest {
.setClustering(CLUSTERING)
.setLabels(LABELS)
.setDecimalTargetTypes(DECIMAL_TARGET_TYPES)
.setConnectionProperties(CONNECTION_PROPERTIES)
.setCreateSession(CREATE_SESSION)
.build();

private static final DatastoreBackupOptions BACKUP_OPTIONS =
Expand Down Expand Up @@ -232,5 +242,7 @@ private void compareLoadConfiguration(
assertEquals(expected.getUseAvroLogicalTypes(), value.getUseAvroLogicalTypes());
assertEquals(expected.getLabels(), value.getLabels());
assertEquals(expected.getDecimalTargetTypes(), value.getDecimalTargetTypes());
assertEquals(expected.getConnectionProperties(), value.getConnectionProperties());
assertEquals(expected.getCreateSession(), value.getCreateSession());
}
}
110 changes: 110 additions & 0 deletions google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java
View file
Expand Up @@ -28,6 +28,7 @@
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;

import com.google.api.client.util.IOUtils;
import com.google.api.gax.paging.Page;
import com.google.auth.oauth2.GoogleCredentials;
import com.google.auth.oauth2.ServiceAccountCredentials;
Expand Down Expand Up @@ -60,6 +61,7 @@
import com.google.cloud.bigquery.ConnectionProperty;
import com.google.cloud.bigquery.ConnectionSettings;
import com.google.cloud.bigquery.CopyJobConfiguration;
import com.google.cloud.bigquery.CsvOptions;
import com.google.cloud.bigquery.Dataset;
import com.google.cloud.bigquery.DatasetId;
import com.google.cloud.bigquery.DatasetInfo;
Expand Down Expand Up @@ -141,9 +143,13 @@
import com.google.gson.JsonObject;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.math.BigDecimal;
import java.nio.ByteBuffer;
import java.nio.channels.Channels;
import java.nio.charset.StandardCharsets;
import java.nio.file.FileSystems;
import java.nio.file.Path;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Time;
Expand Down Expand Up @@ -711,6 +717,36 @@ public class ITBigQueryTest {
private static final List<ConnectionProperty> CONNECTION_PROPERTIES =
ImmutableList.of(CONNECTION_PROPERTY);

private static final Field ID_SCHEMA =
Field.newBuilder("id", LegacySQLTypeName.STRING)
.setMode(Mode.REQUIRED)
.setDescription("id")
.build();
private static final Field FIRST_NAME_SCHEMA =
Field.newBuilder("firstname", LegacySQLTypeName.STRING)
.setMode(Field.Mode.NULLABLE)
.setDescription("First Name")
.build();
private static final Field LAST_NAME_SCHEMA =
Field.newBuilder("lastname", LegacySQLTypeName.STRING)
.setMode(Field.Mode.NULLABLE)
.setDescription("LAST NAME")
.build();
private static final Field EMAIL_SCHEMA =
Field.newBuilder("email", LegacySQLTypeName.STRING)
.setMode(Field.Mode.NULLABLE)
.setDescription("email")
.build();
private static final Field PROFESSION_SCHEMA =
Field.newBuilder("profession", LegacySQLTypeName.STRING)
.setMode(Field.Mode.NULLABLE)
.setDescription("profession")
.build();
private static final Schema SESSION_TABLE_SCHEMA =
Schema.of(ID_SCHEMA, FIRST_NAME_SCHEMA, LAST_NAME_SCHEMA, EMAIL_SCHEMA, PROFESSION_SCHEMA);
private static final Path csvPath =
FileSystems.getDefault().getPath("src/test/resources", "sessionTest.csv").toAbsolutePath();

private static final Set<String> PUBLIC_DATASETS =
ImmutableSet.of("github_repos", "hacker_news", "noaa_gsod", "samples", "usa_names");

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

@Test
public void testLoadSessionSupportWriteChannelConfiguration() throws InterruptedException {
TableId sessionTableId = TableId.of("_SESSION", "test_temp_destination_table_from_file");

WriteChannelConfiguration configuration =
WriteChannelConfiguration.newBuilder(sessionTableId)
.setFormatOptions(CsvOptions.newBuilder().setFieldDelimiter(",").build())
.setCreateDisposition(JobInfo.CreateDisposition.CREATE_IF_NEEDED)
.setSchema(SESSION_TABLE_SCHEMA)
.setCreateSession(true)
.build();
String jobName = "jobId_" + UUID.randomUUID().toString();
JobId jobId = JobId.newBuilder().setLocation("us").setJob(jobName).build();
String sessionId;

// Imports a local file into a table.
try (TableDataWriteChannel writer = bigquery.writer(jobId, configuration);
OutputStream stream = Channels.newOutputStream(writer)) {
InputStream inputStream =
ITBigQueryTest.class.getClassLoader().getResourceAsStream("sessionTest.csv");
// Can use `Files.copy(csvPath, stream);` instead.
// Using IOUtils here because graalvm can't handle resource files.
IOUtils.copy(inputStream, stream);

} catch (IOException e) {
throw new RuntimeException(e);
}
Job loadJob = bigquery.getJob(jobId);
Job completedJob = loadJob.waitFor();

assertNotNull(completedJob);
assertEquals(jobId.getJob(), completedJob.getJobId().getJob());
JobStatistics.LoadStatistics statistics = completedJob.getStatistics();

sessionId = statistics.getSessionInfo().getSessionId();
assertNotNull(sessionId);

// Load job in the same session.
// Should load the data to a temp table.
ConnectionProperty sessionConnectionProperty =
ConnectionProperty.newBuilder().setKey("session_id").setValue(sessionId).build();
WriteChannelConfiguration sessionConfiguration =
WriteChannelConfiguration.newBuilder(sessionTableId)
.setConnectionProperties(ImmutableList.of(sessionConnectionProperty))
.setFormatOptions(CsvOptions.newBuilder().setFieldDelimiter(",").build())
.setCreateDisposition(JobInfo.CreateDisposition.CREATE_IF_NEEDED)
.setSchema(SESSION_TABLE_SCHEMA)
.build();
String sessionJobName = "jobId_" + UUID.randomUUID().toString();
JobId sessionJobId = JobId.newBuilder().setLocation("us").setJob(sessionJobName).build();
try (TableDataWriteChannel writer = bigquery.writer(sessionJobId, sessionConfiguration);
OutputStream stream = Channels.newOutputStream(writer)) {
InputStream inputStream =
ITBigQueryTest.class.getClassLoader().getResourceAsStream("sessionTest.csv");
IOUtils.copy(inputStream, stream);
} catch (IOException e) {
throw new RuntimeException(e);
}
Job queryJobWithSession = bigquery.getJob(sessionJobId);
queryJobWithSession = queryJobWithSession.waitFor();
LoadStatistics statisticsWithSession = queryJobWithSession.getStatistics();
assertNotNull(statisticsWithSession.getSessionInfo().getSessionId());

// Checking if the data loaded to the temp table in the session
String queryTempTable = "SELECT * FROM _SESSION.test_temp_destination_table_from_file;";
QueryJobConfiguration queryJobConfigurationWithSession =
QueryJobConfiguration.newBuilder(queryTempTable)
.setConnectionProperties(ImmutableList.of(sessionConnectionProperty))
.build();
Job queryTempTableJob = bigquery.create(JobInfo.of(queryJobConfigurationWithSession));
queryTempTableJob = queryTempTableJob.waitFor();
assertNotNull(queryTempTableJob.getQueryResults());
}

@Test
public void testLoadSessionSupport() throws InterruptedException {
// Start the session
Expand Down
3 changes: 2 additions & 1 deletion google-cloud-bigquery/src/test/resources/META-INF/native-image/resource-config.json
View file
@@ -1,3 +1,4 @@
{
"resources":[{"pattern": ".*.csv"}]
"resources":[{"pattern": ".*.csv"},
{"pattern": ".*src/test/resources/sessionTest.csv"}]
}
51 changes: 51 additions & 0 deletions google-cloud-bigquery/src/test/resources/sessionTest.csv
View file
@@ -0,0 +1,51 @@
id,firstname,lastname,email,profession
100,Rani,Merell,Rani.Merell@yopmail.com,firefighter
101,Goldie,Dex,Goldie.Dex@yopmail.com,developer
102,Cristabel,Munn,Cristabel.Munn@yopmail.com,developer
103,Genevra,Strephon,Genevra.Strephon@yopmail.com,firefighter
104,Augustine,Thema,Augustine.Thema@yopmail.com,doctor
105,Jemie,Gombach,Jemie.Gombach@yopmail.com,police officer
106,Maye,Stuart,Maye.Stuart@yopmail.com,developer
107,Ayn,Carmena,Ayn.Carmena@yopmail.com,worker
108,Gale,Celestine,Gale.Celestine@yopmail.com,doctor
109,Alex,Jerold,Alex.Jerold@yopmail.com,firefighter
110,Violet,Giule,Violet.Giule@yopmail.com,firefighter
111,Starla,Uird,Starla.Uird@yopmail.com,doctor
112,Tarra,Pelagias,Tarra.Pelagias@yopmail.com,police officer
113,Eugine,Deny,Eugine.Deny@yopmail.com,doctor
114,Shirlee,Ricarda,Shirlee.Ricarda@yopmail.com,doctor
115,Ariela,Penelopa,Ariela.Penelopa@yopmail.com,worker
116,Lelah,Astra,Lelah.Astra@yopmail.com,police officer
117,Debee,Deegan,Debee.Deegan@yopmail.com,developer
118,Pollyanna,Euridice,Pollyanna.Euridice@yopmail.com,worker
119,Cathie,Halsey,Cathie.Halsey@yopmail.com,firefighter
120,Rebeca,Quinn,Rebeca.Quinn@yopmail.com,doctor
121,Paulita,Arquit,Paulita.Arquit@yopmail.com,police officer
122,Rebeca,Emanuel,Rebeca.Emanuel@yopmail.com,firefighter
123,Tera,Ilka,Tera.Ilka@yopmail.com,firefighter
124,Orsola,Briney,Orsola.Briney@yopmail.com,doctor
125,Paulita,Wyn,Paulita.Wyn@yopmail.com,doctor
126,Constance,Christine,Constance.Christine@yopmail.com,firefighter
127,Claresta,Kinnard,Claresta.Kinnard@yopmail.com,developer
128,Leanna,Mendez,Leanna.Mendez@yopmail.com,developer
129,Corina,Chabot,Corina.Chabot@yopmail.com,developer
130,Romona,Audly,Romona.Audly@yopmail.com,worker
131,Cordi,Lynn,Cordi.Lynn@yopmail.com,firefighter
132,Sheree,Tyson,Sheree.Tyson@yopmail.com,worker
133,Jinny,Bevin,Jinny.Bevin@yopmail.com,police officer
134,Kassey,Havens,Kassey.Havens@yopmail.com,firefighter
135,Wanda,Thema,Wanda.Thema@yopmail.com,developer
136,Vita,Jagir,Vita.Jagir@yopmail.com,developer
137,Alie,Aprile,Alie.Aprile@yopmail.com,firefighter
138,Modestia,Jena,Modestia.Jena@yopmail.com,doctor
139,Cyndie,Pelagias,Cyndie.Pelagias@yopmail.com,worker
140,Ariela,Lilybelle,Ariela.Lilybelle@yopmail.com,firefighter
141,Jan,Parette,Jan.Parette@yopmail.com,firefighter
142,Merry,Horan,Merry.Horan@yopmail.com,developer
143,Katuscha,Candy,Katuscha.Candy@yopmail.com,police officer
144,Kerrin,Heisel,Kerrin.Heisel@yopmail.com,developer
145,Nollie,Magdalen,Nollie.Magdalen@yopmail.com,doctor
146,Karlee,Gordon,Karlee.Gordon@yopmail.com,developer
147,Dolli,Fadiman,Dolli.Fadiman@yopmail.com,firefighter
148,Leontine,Delp,Leontine.Delp@yopmail.com,worker
149,Ricky,Nadia,Ricky.Nadia@yopmail.com,doctor

0 comments on commit 42851d8

Please sign in to comment.