From 09e12f0c8070ba3eef681248b776c97bbba82a28 Mon Sep 17 00:00:00 2001 From: Phong Chuong <147636638+PhongChuong@users.noreply.github.com> Date: Fri, 17 Nov 2023 13:56:06 -0500 Subject: [PATCH 1/3] feat: Add Routine DataGovernanceType This also: - fix an issue where RoutineInfo unit test was being skipped - fix an issue where ImportLibaries were not being exported in the toPb method --- .../com/google/cloud/bigquery/Routine.java | 6 ++++ .../google/cloud/bigquery/RoutineInfo.java | 35 +++++++++++++++++-- .../cloud/bigquery/RoutineInfoTest.java | 7 ++++ .../google/cloud/bigquery/RoutineTest.java | 5 +++ .../cloud/bigquery/it/ITBigQueryTest.java | 27 ++++++++++++++ 5 files changed, 78 insertions(+), 2 deletions(-) diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Routine.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Routine.java index d2d6ec683..5d99ccfbd 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Routine.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Routine.java @@ -135,6 +135,12 @@ public Builder setRemoteFunctionOptions(RemoteFunctionOptions remoteFunctionOpti return this; } + @Override + public Builder setDataGovernanceType(String dataGovernanceType) { + infoBuilder.setDataGovernanceType(dataGovernanceType); + return this; + } + @Override public Routine build() { return new Routine(bigquery, infoBuilder); diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/RoutineInfo.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/RoutineInfo.java index 3001d2fb9..7cea99f93 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/RoutineInfo.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/RoutineInfo.java @@ -72,6 +72,8 @@ public Routine apply(RoutineInfo routineInfo) { private final String body; private final RemoteFunctionOptions remoteFunctionOptions; + private final String dataGovernanceType; + public abstract static class Builder { abstract Builder setRoutineId(RoutineId id); @@ -157,6 +159,13 @@ public abstract static class Builder { */ public abstract Builder setRemoteFunctionOptions(RemoteFunctionOptions remoteFunctionOptions); + /** + * Sets the data governance type for the Builder (e.g. DATA_MASKING). + * + *
See https://cloud.google.com/bigquery/docs/reference/rest/v2/routines
+ */
+ public abstract Builder setDataGovernanceType(String dataGovernanceType);
+
/** Creates a {@code RoutineInfo} object. */
public abstract RoutineInfo build();
}
@@ -177,6 +186,8 @@ static class BuilderImpl extends Builder {
private String body;
private RemoteFunctionOptions remoteFunctionOptions;
+ private String dataGovernanceType;
+
BuilderImpl() {}
BuilderImpl(RoutineInfo routineInfo) {
@@ -194,6 +205,7 @@ static class BuilderImpl extends Builder {
this.importedLibrariesList = routineInfo.importedLibrariesList;
this.body = routineInfo.body;
this.remoteFunctionOptions = routineInfo.remoteFunctionOptions;
+ this.dataGovernanceType = routineInfo.dataGovernanceType;
}
BuilderImpl(Routine routinePb) {
@@ -225,6 +237,7 @@ static class BuilderImpl extends Builder {
this.remoteFunctionOptions =
RemoteFunctionOptions.fromPb(routinePb.getRemoteFunctionOptions());
}
+ this.dataGovernanceType = routinePb.getDataGovernanceType();
}
@Override
@@ -311,6 +324,12 @@ public Builder setRemoteFunctionOptions(RemoteFunctionOptions remoteFunctionOpti
return this;
}
+ @Override
+ public Builder setDataGovernanceType(String dataGovernanceType) {
+ this.dataGovernanceType = dataGovernanceType;
+ return this;
+ }
+
@Override
public RoutineInfo build() {
return new RoutineInfo(this);
@@ -332,6 +351,7 @@ public RoutineInfo build() {
this.importedLibrariesList = builder.importedLibrariesList;
this.body = builder.body;
this.remoteFunctionOptions = builder.remoteFunctionOptions;
+ this.dataGovernanceType = builder.dataGovernanceType;
}
/** Returns the RoutineId identified for the routine resource. * */
@@ -411,6 +431,11 @@ public RemoteFunctionOptions getRemoteFunctionOptions() {
return remoteFunctionOptions;
};
+ /** Returns the data governance type of the routine, e.g. DATA_MASKING. */
+ public String getDataGovernanceType() {
+ return dataGovernanceType;
+ }
+
/** Returns a builder pre-populated using the current values of this routine. */
public Builder toBuilder() {
return new BuilderImpl(this);
@@ -433,6 +458,7 @@ public String toString() {
.add("importedLibrariesList", importedLibrariesList)
.add("body", body)
.add("remoteFunctionOptions", remoteFunctionOptions)
+ .add("dataGovernanceType", dataGovernanceType)
.toString();
}
@@ -452,7 +478,8 @@ public int hashCode() {
returnTableType,
importedLibrariesList,
body,
- remoteFunctionOptions);
+ remoteFunctionOptions,
+ dataGovernanceType);
}
@Override
@@ -490,7 +517,8 @@ Routine toPb() {
.setDescription(getDescription())
.setDeterminismLevel(getDeterminismLevel())
.setLastModifiedTime(getLastModifiedTime())
- .setLanguage(getLanguage());
+ .setLanguage(getLanguage())
+ .setDataGovernanceType(getDataGovernanceType());
if (getRoutineId() != null) {
routinePb.setRoutineReference(getRoutineId().toPb());
}
@@ -506,6 +534,9 @@ Routine toPb() {
if (getRemoteFunctionOptions() != null) {
routinePb.setRemoteFunctionOptions(getRemoteFunctionOptions().toPb());
}
+ if (getImportedLibraries() != null) {
+ routinePb.setImportedLibraries(getImportedLibraries());
+ }
return routinePb;
}
diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/RoutineInfoTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/RoutineInfoTest.java
index 1f1181433..ae061b62f 100644
--- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/RoutineInfoTest.java
+++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/RoutineInfoTest.java
@@ -33,6 +33,8 @@ public class RoutineInfoTest {
private static final Long LAST_MODIFIED_TIME = 20L;
private static final String LANGUAGE = "SQL";
+ private static final String DATA_GOVERNANCE_TYPE = "DATA_MASKING";
+
private static final RoutineArgument ARG_1 =
RoutineArgument.newBuilder()
.setDataType(StandardSQLDataType.newBuilder("STRING").build())
@@ -63,6 +65,7 @@ public class RoutineInfoTest {
.setReturnType(RETURN_TYPE)
.setImportedLibraries(IMPORTED_LIBRARIES)
.setBody(BODY)
+ .setDataGovernanceType(DATA_GOVERNANCE_TYPE)
.build();
@Test
@@ -90,6 +93,7 @@ public void testBuilder() {
assertEquals(RETURN_TYPE, ROUTINE_INFO.getReturnType());
assertEquals(IMPORTED_LIBRARIES, ROUTINE_INFO.getImportedLibraries());
assertEquals(BODY, ROUTINE_INFO.getBody());
+ assertEquals(DATA_GOVERNANCE_TYPE, ROUTINE_INFO.getDataGovernanceType());
}
@Test
@@ -107,8 +111,10 @@ public void testOf() {
assertNull(routineInfo.getReturnType());
assertNull(routineInfo.getImportedLibraries());
assertNull(routineInfo.getBody());
+ assertNull(routineInfo.getDataGovernanceType());
}
+ @Test
public void testToAndFromPb() {
compareRoutineInfo(ROUTINE_INFO, RoutineInfo.fromPb(ROUTINE_INFO.toPb()));
}
@@ -132,6 +138,7 @@ public void compareRoutineInfo(RoutineInfo expected, RoutineInfo value) {
assertEquals(expected.getReturnType(), value.getReturnType());
assertEquals(expected.getImportedLibraries(), value.getImportedLibraries());
assertEquals(expected.getBody(), value.getBody());
+ assertEquals(expected.getDataGovernanceType(), value.getDataGovernanceType());
assertEquals(expected.hashCode(), value.hashCode());
assertEquals(expected.toString(), value.toString());
}
diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/RoutineTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/RoutineTest.java
index c9080e851..eaf142012 100644
--- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/RoutineTest.java
+++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/RoutineTest.java
@@ -91,6 +91,8 @@ public class RoutineTest {
.setMaxBatchingRows(10L)
.build();
+ private static final String DATA_GOVERNANCE_TYPE = "DATA_MASKING";
+
private static final RoutineInfo ROUTINE_INFO =
RoutineInfo.newBuilder(ROUTINE_ID)
.setEtag(ETAG)
@@ -104,6 +106,7 @@ public class RoutineTest {
.setImportedLibraries(IMPORTED_LIBRARIES)
.setBody(BODY)
.setRemoteFunctionOptions(REMOTE_FUNCTION_OPTIONS)
+ .setDataGovernanceType(DATA_GOVERNANCE_TYPE)
.build();
private static final RoutineInfo ROUTINE_INFO_TVF =
@@ -146,6 +149,7 @@ public void testBuilder() {
.setImportedLibraries(IMPORTED_LIBRARIES)
.setBody(BODY)
.setRemoteFunctionOptions(REMOTE_FUNCTION_OPTIONS)
+ .setDataGovernanceType(DATA_GOVERNANCE_TYPE)
.build();
assertEquals(ETAG, builtRoutine.getEtag());
assertEquals(DETERMINISM_LEVEL, builtRoutine.getDeterminismLevel());
@@ -247,5 +251,6 @@ public void compareRoutineInfo(RoutineInfo expected, RoutineInfo value) {
assertEquals(expected.getBody(), value.getBody());
assertEquals(expected.hashCode(), value.hashCode());
assertEquals(expected.getRemoteFunctionOptions(), value.getRemoteFunctionOptions());
+ assertEquals(expected.getDataGovernanceType(), value.getDataGovernanceType());
}
}
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 b37f389ae..7bdb7a7fb 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
@@ -2684,6 +2684,33 @@ public void testRoutineAPICreationTVF() {
assertEquals(routine.getReturnTableType(), returnTableType);
}
+ @Test
+ public void testRoutineDataGovernanceType() {
+ String routineName = RemoteBigQueryHelper.generateRoutineName();
+ RoutineId routineId = RoutineId.of(ROUTINE_DATASET, routineName);
+ RoutineInfo routineInfo =
+ RoutineInfo.newBuilder(routineId)
+ .setLanguage("SQL")
+ .setRoutineType("SCALAR_FUNCTION")
+ .setBody("x")
+ .setArguments(
+ ImmutableList.of(
+ RoutineArgument.newBuilder()
+ .setName("x")
+ .setDataType(StandardSQLDataType.newBuilder("INT64").build())
+ .build()))
+ .setReturnType(StandardSQLDataType.newBuilder("INT64").build())
+ .setDataGovernanceType("DATA_MASKING")
+ .build();
+
+ Routine routine = bigquery.create(routineInfo);
+ assertNotNull(routine);
+ assertEquals(routine.getLanguage(), "SQL");
+ assertEquals(routine.getRoutineType(), "SCALAR_FUNCTION");
+ assertEquals(routine.getReturnType(), StandardSQLDataType.newBuilder("INT64").build());
+ assertEquals(routine.getDataGovernanceType(), "DATA_MASKING");
+ }
+
@Test
public void testAuthorizeRoutine() {
String routineName = RemoteBigQueryHelper.generateRoutineName();
From 92b33b91a189a97d7b923dbfebbfa8d1d20b1d2b Mon Sep 17 00:00:00 2001
From: Owl Bot