diff --git a/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Model.java b/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Model.java index 22dcfa0db527..abe1f0f2ba77 100644 --- a/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Model.java +++ b/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Model.java @@ -126,6 +126,12 @@ Builder setFeatureColumns(List featureColumnList) { return this; } + @Override + public Builder setEncryptionConfiguration(EncryptionConfiguration configuration) { + infoBuilder.setEncryptionConfiguration(configuration); + return this; + } + public Model build() { return new Model(bigquery, infoBuilder); } diff --git a/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ModelInfo.java b/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ModelInfo.java index d78bdbfbafab..5796c820f306 100644 --- a/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ModelInfo.java +++ b/google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ModelInfo.java @@ -71,6 +71,7 @@ public Model apply(ModelInfo ModelInfo) { private final ImmutableList trainingRunList; private final ImmutableList featureColumnList; private final ImmutableList labelColumnList; + private final EncryptionConfiguration encryptionConfiguration; /** A builder for {@code ModelInfo} objects. */ public abstract static class Builder { @@ -112,6 +113,8 @@ public abstract static class Builder { abstract Builder setFeatureColumns(List featureColumnList); + public abstract Builder setEncryptionConfiguration(EncryptionConfiguration configuration); + /** Creates a {@code ModelInfo} object. */ public abstract ModelInfo build(); } @@ -130,6 +133,7 @@ static class BuilderImpl extends Builder { private List trainingRunList = Collections.emptyList(); private List labelColumnList = Collections.emptyList(); private List featureColumnList = Collections.emptyList(); + private EncryptionConfiguration encryptionConfiguration; BuilderImpl() {} @@ -145,6 +149,7 @@ static class BuilderImpl extends Builder { this.trainingRunList = modelInfo.trainingRunList; this.labelColumnList = modelInfo.labelColumnList; this.featureColumnList = modelInfo.featureColumnList; + this.encryptionConfiguration = modelInfo.encryptionConfiguration; } BuilderImpl(Model modelPb) { @@ -171,6 +176,10 @@ static class BuilderImpl extends Builder { this.featureColumnList = Lists.transform(modelPb.getFeatureColumns(), StandardSQLField.FROM_PB_FUNCTION); } + if (modelPb.getEncryptionConfiguration() != null) { + this.encryptionConfiguration = + new EncryptionConfiguration.Builder(modelPb.getEncryptionConfiguration()).build(); + } } @Override @@ -245,6 +254,12 @@ Builder setFeatureColumns(List featureColumnList) { return this; } + @Override + public Builder setEncryptionConfiguration(EncryptionConfiguration configuration) { + this.encryptionConfiguration = configuration; + return this; + } + @Override public ModelInfo build() { return new ModelInfo(this); @@ -264,6 +279,7 @@ public ModelInfo build() { this.trainingRunList = ImmutableList.copyOf(builder.trainingRunList); this.labelColumnList = ImmutableList.copyOf(builder.labelColumnList); this.featureColumnList = ImmutableList.copyOf(builder.featureColumnList); + this.encryptionConfiguration = builder.encryptionConfiguration; } /** Returns the hash of the model resource. */ @@ -332,6 +348,10 @@ public ImmutableList getFeatureColumns() { return featureColumnList; } + public EncryptionConfiguration getEncryptionConfiguration() { + return encryptionConfiguration; + } + public Builder toBuilder() { return new BuilderImpl(this); } @@ -351,6 +371,7 @@ public String toString() { .add("trainingRuns", trainingRunList) .add("labelColumns", labelColumnList) .add("featureColumns", featureColumnList) + .add("encryptionConfiguration", encryptionConfiguration) .toString(); } @@ -403,6 +424,9 @@ Model toPb() { modelPb.setFeatureColumns( Lists.transform(featureColumnList, StandardSQLField.TO_PB_FUNCTION)); } + if (encryptionConfiguration != null) { + modelPb.setEncryptionConfiguration(encryptionConfiguration.toPb()); + } return modelPb; } diff --git a/google-cloud-clients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ModelInfoTest.java b/google-cloud-clients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ModelInfoTest.java index 891ba07ff084..2657ccc4497e 100644 --- a/google-cloud-clients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ModelInfoTest.java +++ b/google-cloud-clients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ModelInfoTest.java @@ -33,6 +33,8 @@ public class ModelInfoTest { private static final Long EXPIRATION_TIME = 30L; private static final String DESCRIPTION = "description"; private static final String FRIENDLY_NAME = "friendlyname"; + private static final EncryptionConfiguration MODEL_ENCRYPTION_CONFIGURATION = + EncryptionConfiguration.newBuilder().setKmsKeyName("KMS_KEY_1").build(); private static final TrainingOptions TRAINING_OPTIONS = new TrainingOptions().setDataSplitColumn("foo").setEarlyStop(true).setLossType("bar"); @@ -49,6 +51,7 @@ public class ModelInfoTest { .setDescription(DESCRIPTION) .setFriendlyName(FRIENDLY_NAME) .setTrainingRuns(TRAINING_RUN_LIST) + .setEncryptionConfiguration(MODEL_ENCRYPTION_CONFIGURATION) .build(); @Test @@ -71,6 +74,7 @@ public void testBuilder() { assertEquals(DESCRIPTION, MODEL_INFO.getDescription()); assertEquals(FRIENDLY_NAME, MODEL_INFO.getFriendlyName()); assertEquals(TRAINING_OPTIONS, MODEL_INFO.getTrainingRuns().get(0).getTrainingOptions()); + assertEquals(MODEL_ENCRYPTION_CONFIGURATION, MODEL_INFO.getEncryptionConfiguration()); } @Test @@ -83,6 +87,7 @@ public void testOf() { assertNull(modelInfo.getExpirationTime()); assertNull(modelInfo.getDescription()); assertNull(modelInfo.getFriendlyName()); + assertNull(modelInfo.getEncryptionConfiguration()); assertEquals(modelInfo.getTrainingRuns().isEmpty(), true); assertEquals(modelInfo.getLabelColumns().isEmpty(), true); assertEquals(modelInfo.getFeatureColumns().isEmpty(), true); @@ -112,5 +117,6 @@ private void compareModelInfo(ModelInfo expected, ModelInfo value) { assertEquals(expected.getTrainingRuns(), value.getTrainingRuns()); assertEquals(expected.getLabelColumns(), value.getLabelColumns()); assertEquals(expected.getFeatureColumns(), value.getFeatureColumns()); + assertEquals(expected.getEncryptionConfiguration(), value.getEncryptionConfiguration()); } }