bigquery: add support for CMEK functionality on ML models.

This change exposes the CMEK EncryptionConfig for ML models, which
mirrors the functionality for tables and allows users to define
customer-managed encryption keys for BigQuery storage related to
ML models.

Fixes: #1579

Change-Id: I523501124a13a134f4b14701e9d59c0ce4072683
Reviewed-on: https://code-review.googlesource.com/c/gocloud/+/45390
Reviewed-by: kokoro <noreply+kokoro@google.com>
Reviewed-by: Tyler Bui-Palsulich <tbp@google.com>
Reviewed-by: Alex Hong <hongalex@google.com>

bigquery/model.go

@@ -125,6 +125,9 @@ type ModelMetadata struct {
 	// inherited from the encapsulating dataset.
 	Location string
 
+	// Custom encryption configuration (e.g., Cloud KMS keys).
+	EncryptionConfig *EncryptionConfig
+
 	// The input feature columns used to train the model.
 	featureColumns []*bq.StandardSqlField
 
@@ -201,6 +204,7 @@ func bqToModelMetadata(m *bq.Model) (*ModelMetadata, error) {
 		ExpirationTime:   unixMillisToTime(m.ExpirationTime),
 		CreationTime:     unixMillisToTime(m.CreationTime),
 		LastModifiedTime: unixMillisToTime(m.LastModifiedTime),
+		EncryptionConfig: bqToEncryptionConfig(m.EncryptionConfiguration),
 		featureColumns:   m.FeatureColumns,
 		labelColumns:     m.LabelColumns,
 		trainingRuns:     m.TrainingRuns,
@@ -222,6 +226,9 @@ type ModelMetadataToUpdate struct {
 	// set ExpirationTime to NeverExpire.  The zero value is ignored.
 	ExpirationTime time.Time
 
+	// The model's encryption configuration.
+	EncryptionConfig *EncryptionConfig
+
 	labelUpdater
 }
 
@@ -241,6 +248,10 @@ func (mm *ModelMetadataToUpdate) toBQ() (*bq.Model, error) {
 		forceSend("FriendlyName")
 	}
 
+	if mm.EncryptionConfig != nil {
+		m.EncryptionConfiguration = mm.EncryptionConfig.toBQ()
+	}
+
 	if !validExpiration(mm.ExpirationTime) {
 		return nil, invalidTimeError(mm.ExpirationTime)
 	}

bigquery/model_test.go

@@ -33,14 +33,15 @@ func TestBQToModelMetadata(t *testing.T) {
 		{&bq.Model{}, &ModelMetadata{}},
 		{
 			&bq.Model{
-				CreationTime:     aTimeMillis,
-				Description:      "desc",
-				Etag:             "etag",
-				ExpirationTime:   aTimeMillis,
-				FriendlyName:     "fname",
-				LastModifiedTime: aTimeMillis,
-				Location:         "loc",
-				Labels:           map[string]string{"a": "b"},
+				CreationTime:            aTimeMillis,
+				Description:             "desc",
+				Etag:                    "etag",
+				ExpirationTime:          aTimeMillis,
+				EncryptionConfiguration: &bq.EncryptionConfiguration{KmsKeyName: "keyName"},
+				FriendlyName:            "fname",
+				LastModifiedTime:        aTimeMillis,
+				Location:                "loc",
+				Labels:                  map[string]string{"a": "b"},
 			},
 			&ModelMetadata{
 				CreationTime:     aTime.Truncate(time.Millisecond),
@@ -49,6 +50,7 @@ func TestBQToModelMetadata(t *testing.T) {
 				ExpirationTime:   aTime.Truncate(time.Millisecond),
 				Name:             "fname",
 				LastModifiedTime: aTime.Truncate(time.Millisecond),
+				EncryptionConfig: &EncryptionConfig{KMSKeyName: "keyName"},
 				Location:         "loc",
 				Labels:           map[string]string{"a": "b"},
 			},

bigquery/table.go

@@ -232,7 +232,7 @@ func bqToClustering(q *bq.Clustering) *Clustering {
 	}
 }
 
-// EncryptionConfig configures customer-managed encryption on tables.
+// EncryptionConfig configures customer-managed encryption on tables and ML models.
 type EncryptionConfig struct {
 	// Describes the Cloud KMS encryption key that will be used to protect
 	// destination BigQuery table. The BigQuery Service Account associated with your
@@ -587,8 +587,7 @@ type TableMetadataToUpdate struct {
 	// When updating a schema, you can add columns but not remove them.
 	Schema Schema
 
-	// The table's encryption configuration.  When calling Update, ensure that
-	// all mutable fields of EncryptionConfig are populated.
+	// The table's encryption configuration.
 	EncryptionConfig *EncryptionConfig
 
 	// The time when this table expires. To remove a table's expiration,