feat: Added encryption_spec_key_name to all relevant resources

google/cloud/aiplatform/datasets/dataset.py

@@ -25,6 +25,7 @@
 from google.cloud.aiplatform import utils
 
 from google.cloud.aiplatform.datasets import _datasources
+from google.cloud.aiplatform_v1beta1.types import encryption_spec as gca_encryption_spec
 from google.cloud.aiplatform_v1beta1.types import io as gca_io
 from google.cloud.aiplatform_v1beta1.types import dataset as gca_dataset
 from google.cloud.aiplatform_v1beta1.services.dataset_service import (
@@ -69,7 +70,9 @@ def __init__(
 
         """
 
-        super().__init__(project=project, location=location, credentials=credentials)
+        super().__init__(
+            project=project, location=location, credentials=credentials,
+        )
         self._gca_resource = self._get_gca_resource(resource_name=dataset_name)
         self._validate_metadata_schema_uri()
 
@@ -106,6 +109,7 @@ def create(
         location: Optional[str] = None,
         credentials: Optional[auth_credentials.Credentials] = None,
         request_metadata: Optional[Sequence[Tuple[str, str]]] = (),
+        encryption_spec_key_name: Optional[str] = None,
         sync: bool = True,
     ) -> "Dataset":
         """Creates a new dataset and optionally imports data into dataset when
@@ -166,6 +170,17 @@ def create(
                 credentials set in aiplatform.init.
             request_metadata (Sequence[Tuple[str, str]]):
                 Strings which should be sent along with the request as metadata.
+            encryption_spec_key_name (Optional[str]):
+                Optional. The Cloud KMS resource identifier of the customer
+                managed encryption key used to protect the dataset. Has the
+                form:
+                ``projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key``.
+                The key needs to be in the same region as where the compute
+                resource is created.
+
+                If set, this Dataset and all sub-resources of this Dataset will be secured by this key.
+
+                Overrides encryption_spec_key_name set in aiplatform.init.
             sync (bool):
                 Whether to execute this method synchronously. If False, this method
                 will be executed in concurrent Future and any downstream object will
@@ -201,6 +216,9 @@ def create(
             location=location or initializer.global_config.location,
             credentials=credentials or initializer.global_config.credentials,
             request_metadata=request_metadata,
+            encryption_spec=initializer.global_config.get_encryption_spec(
+                encryption_spec_key_name=encryption_spec_key_name
+            ),
             sync=sync,
         )
 
@@ -217,6 +235,7 @@ def _create_and_import(
         location: str,
         credentials: Optional[auth_credentials.Credentials],
         request_metadata: Optional[Sequence[Tuple[str, str]]] = (),
+        encryption_spec: Optional[gca_encryption_spec.EncryptionSpec] = None,
         sync: bool = True,
     ) -> "Dataset":
         """Creates a new dataset and optionally imports data into dataset when
@@ -253,6 +272,12 @@ def _create_and_import(
                 credentials set in aiplatform.init.
             request_metadata (Sequence[Tuple[str, str]]):
                 Strings which should be sent along with the request as metadata.
+            encryption_spec (Optional[gca_encryption_spec.EncryptionSpec]):
+                Optional. The Cloud KMS customer managed encryption key used to protect the dataset.
+                The key needs to be in the same region as where the compute
+                resource is created.
+
+                If set, this Dataset and all sub-resources of this Dataset will be secured by this key.
             sync (bool):
                 Whether to execute this method synchronously. If False, this method
                 will be executed in concurrent Future and any downstream object will
@@ -270,6 +295,7 @@ def _create_and_import(
             metadata_schema_uri=metadata_schema_uri,
             datasource=datasource,
             request_metadata=request_metadata,
+            encryption_spec=encryption_spec,
         )
 
         created_dataset = create_dataset_lro.result()
@@ -297,6 +323,7 @@ def _create(
         metadata_schema_uri: str,
         datasource: _datasources.Datasource,
         request_metadata: Sequence[Tuple[str, str]] = (),
+        encryption_spec: Optional[gca_encryption_spec.EncryptionSpec] = None,
     ) -> operation.Operation:
         """Creates a new managed dataset by directly calling API client.
 
@@ -323,15 +350,22 @@ def _create(
             request_metadata (Sequence[Tuple[str, str]]):
                 Strings which should be sent along with the create_dataset
                 request as metadata. Usually to specify special dataset config.
+            encryption_spec (Optional[gca_encryption_spec.EncryptionSpec]):
+                Optional. The Cloud KMS customer managed encryption key used to protect the dataset.
+                The key needs to be in the same region as where the compute
+                resource is created.
 
+                If set, this Dataset and all sub-resources of this Dataset will be secured by this key.
         Returns:
             operation (Operation):
                 An object representing a long-running operation.
         """
+
         gapic_dataset = gca_dataset.Dataset(
             display_name=display_name,
             metadata_schema_uri=metadata_schema_uri,
             metadata=datasource.dataset_metadata,
+            encryption_spec=encryption_spec,
         )
 
         return api_client.create_dataset(

google/cloud/aiplatform/datasets/image_dataset.py

@@ -44,6 +44,7 @@ def create(
         location: Optional[str] = None,
         credentials: Optional[auth_credentials.Credentials] = None,
         request_metadata: Optional[Sequence[Tuple[str, str]]] = (),
+        encryption_spec_key_name: Optional[str] = None,
         sync: bool = True,
     ) -> "ImageDataset":
         """Creates a new image dataset and optionally imports data into dataset when
@@ -94,6 +95,17 @@ def create(
                 credentials set in aiplatform.init.
             request_metadata (Sequence[Tuple[str, str]]):
                 Strings which should be sent along with the request as metadata.
+            encryption_spec_key_name (Optional[str]):
+                Optional. The Cloud KMS resource identifier of the customer
+                managed encryption key used to protect the dataset. Has the
+                form:
+                ``projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key``.
+                The key needs to be in the same region as where the compute
+                resource is created.
+
+                If set, this Dataset and all sub-resources of this Dataset will be secured by this key.
+
+                Overrides encryption_spec_key_name set in aiplatform.init.
             sync (bool):
                 Whether to execute this method synchronously. If False, this method
                 will be executed in concurrent Future and any downstream object will
@@ -130,5 +142,8 @@ def create(
             location=location or initializer.global_config.location,
             credentials=credentials or initializer.global_config.credentials,
             request_metadata=request_metadata,
+            encryption_spec=initializer.global_config.get_encryption_spec(
+                encryption_spec_key_name=encryption_spec_key_name
+            ),
             sync=sync,
         )

google/cloud/aiplatform/datasets/tabular_dataset.py

@@ -43,6 +43,7 @@ def create(
         location: Optional[str] = None,
         credentials: Optional[auth_credentials.Credentials] = None,
         request_metadata: Optional[Sequence[Tuple[str, str]]] = (),
+        encryption_spec_key_name: Optional[str] = None,
         sync: bool = True,
     ) -> "TabularDataset":
         """Creates a new tabular dataset.
@@ -75,6 +76,17 @@ def create(
                 credentials set in aiplatform.init.
             request_metadata (Sequence[Tuple[str, str]]):
                 Strings which should be sent along with the request as metadata.
+            encryption_spec_key_name (Optional[str]):
+                Optional. The Cloud KMS resource identifier of the customer
+                managed encryption key used to protect the dataset. Has the
+                form:
+                ``projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key``.
+                The key needs to be in the same region as where the compute
+                resource is created.
+
+                If set, this Dataset and all sub-resources of this Dataset will be secured by this key.
+
+                Overrides encryption_spec_key_name set in aiplatform.init.
             sync (bool):
                 Whether to execute this method synchronously. If False, this method
                 will be executed in concurrent Future and any downstream object will
@@ -110,6 +122,9 @@ def create(
             location=location or initializer.global_config.location,
             credentials=credentials or initializer.global_config.credentials,
             request_metadata=request_metadata,
+            encryption_spec=initializer.global_config.get_encryption_spec(
+                encryption_spec_key_name=encryption_spec_key_name
+            ),
             sync=sync,
         )
 

google/cloud/aiplatform/datasets/text_dataset.py

@@ -44,6 +44,7 @@ def create(
         location: Optional[str] = None,
         credentials: Optional[auth_credentials.Credentials] = None,
         request_metadata: Optional[Sequence[Tuple[str, str]]] = (),
+        encryption_spec_key_name: Optional[str] = None,
         sync: bool = True,
     ) -> "TextDataset":
         """Creates a new text dataset and optionally imports data into dataset when
@@ -101,6 +102,17 @@ def create(
                 credentials set in aiplatform.init.
             request_metadata (Sequence[Tuple[str, str]]):
                 Strings which should be sent along with the request as metadata.
+            encryption_spec_key_name (Optional[str]):
+                Optional. The Cloud KMS resource identifier of the customer
+                managed encryption key used to protect the dataset. Has the
+                form:
+                ``projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key``.
+                The key needs to be in the same region as where the compute
+                resource is created.
+
+                If set, this Dataset and all sub-resources of this Dataset will be secured by this key.
+
+                Overrides encryption_spec_key_name set in aiplatform.init.
             sync (bool):
                 Whether to execute this method synchronously. If False, this method
                 will be executed in concurrent Future and any downstream object will
@@ -137,5 +149,8 @@ def create(
             location=location or initializer.global_config.location,
             credentials=credentials or initializer.global_config.credentials,
             request_metadata=request_metadata,
+            encryption_spec=initializer.global_config.get_encryption_spec(
+                encryption_spec_key_name=encryption_spec_key_name
+            ),
             sync=sync,
         )

google/cloud/aiplatform/datasets/video_dataset.py

@@ -44,6 +44,7 @@ def create(
         location: Optional[str] = None,
         credentials: Optional[auth_credentials.Credentials] = None,
         request_metadata: Optional[Sequence[Tuple[str, str]]] = (),
+        encryption_spec_key_name: Optional[str] = None,
         sync: bool = True,
     ) -> "VideoDataset":
         """Creates a new video dataset and optionally imports data into dataset when
@@ -94,6 +95,17 @@ def create(
                 credentials set in aiplatform.init.
             request_metadata (Sequence[Tuple[str, str]]):
                 Strings which should be sent along with the request as metadata.
+            encryption_spec_key_name (Optional[str]):
+                Optional. The Cloud KMS resource identifier of the customer
+                managed encryption key used to protect the dataset. Has the
+                form:
+                ``projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key``.
+                The key needs to be in the same region as where the compute
+                resource is created.
+
+                If set, this Dataset and all sub-resources of this Dataset will be secured by this key.
+
+                Overrides encryption_spec_key_name set in aiplatform.init.
             sync (bool):
                 Whether to execute this method synchronously. If False, this method
                 will be executed in concurrent Future and any downstream object will
@@ -130,5 +142,8 @@ def create(
             location=location or initializer.global_config.location,
             credentials=credentials or initializer.global_config.credentials,
             request_metadata=request_metadata,
+            encryption_spec=initializer.global_config.get_encryption_spec(
+                encryption_spec_key_name=encryption_spec_key_name
+            ),
             sync=sync,
         )

google/cloud/aiplatform/initializer.py

@@ -29,6 +29,7 @@
 from google.auth.exceptions import GoogleAuthError
 from google.cloud.aiplatform import utils
 from google.cloud.aiplatform import constants
+from google.cloud.aiplatform_v1beta1.types import encryption_spec as gca_encryption_spec
 
 
 class _Config:
@@ -40,6 +41,7 @@ def __init__(self):
         self._location = None
         self._staging_bucket = None
         self._credentials = None
+        self._encryption_spec_key_name = None
 
     def init(
         self,
@@ -49,6 +51,7 @@ def init(
         experiment: Optional[str] = None,
         staging_bucket: Optional[str] = None,
         credentials: Optional[auth_credentials.Credentials] = None,
+        encryption_spec_key_name: Optional[str] = None,
     ):
         """Updates common initalization parameters with provided options.
 
@@ -59,9 +62,18 @@ def init(
             experiment (str): The experiment to assign
             staging_bucket (str): The default staging bucket to use to stage artifacts
                 when making API calls. In the form gs://...
-            credentials (google.auth.crendentials.Crendentials): The default custom
+            credentials (google.auth.crendentials.Credentials): The default custom
                 credentials to use when making API calls. If not provided crendentials
                 will be ascertained from the environment.
+            encryption_spec_key_name (Optional[str]):
+                Optional. The Cloud KMS resource identifier of the customer
+                managed encryption key used to protect a resource. Has the
+                form:
+                ``projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key``.
+                The key needs to be in the same region as where the compute
+                resource is created.
+
+                If set, this resource and all sub-resources will be secured by this key.
         """
         if project:
             self._project = project
@@ -75,6 +87,25 @@ def init(
             self._staging_bucket = staging_bucket
         if credentials:
             self._credentials = credentials
+        if encryption_spec_key_name:
+            self._encryption_spec_key_name = encryption_spec_key_name
+
+    def get_encryption_spec(
+        self, encryption_spec_key_name: Optional[str]
+    ) -> Optional[gca_encryption_spec.EncryptionSpec]:
+        """Creates a gca_encryption_spec.EncryptionSpec instance from the given key name.
+        If the provided key name is None, it uses the default key name if provided.
+
+        Args:
+            encryption_spec_key_name (Optional[str]): The default encryption key name to use when creating resources.
+        """
+        kms_key_name = encryption_spec_key_name or self.encryption_spec_key_name
+        encryption_spec = None
+        if kms_key_name:
+            encryption_spec = gca_encryption_spec.EncryptionSpec(
+                kms_key_name=kms_key_name
+            )
+        return encryption_spec
 
     @property
     def project(self) -> str:
@@ -126,6 +157,11 @@ def credentials(self) -> Optional[auth_credentials.Credentials]:
         logger.removeFilter(logging_warning_filter)
         return credentials
 
+    @property
+    def encryption_spec_key_name(self) -> Optional[str]:
+        """Default encryption spec key name, if provided."""
+        return self._encryption_spec_key_name
+
     def get_client_options(
         self, location_override: Optional[str] = None, prediction_client: bool = False,
     ) -> client_options.ClientOptions:

google/cloud/aiplatform/jobs.py

@@ -42,6 +42,7 @@
 from google.cloud.aiplatform_v1beta1.types import (
     machine_resources as gca_machine_resources,
 )
+
 from google.cloud.aiplatform_v1beta1.types import explanation as gca_explanation
 
 logging.basicConfig(level=logging.INFO, stream=sys.stdout)
@@ -238,6 +239,7 @@ def create(
         project: Optional[str] = None,
         location: Optional[str] = None,
         credentials: Optional[auth_credentials.Credentials] = None,
+        encryption_spec_key_name: Optional[str] = None,
         sync: bool = True,
     ) -> "BatchPredictionJob":
         """Create a batch prediction job.
@@ -372,6 +374,19 @@ def create(
             credentials (Optional[auth_credentials.Credentials]):
                 Custom credentials to use to create this batch prediction
                 job. Overrides credentials set in aiplatform.init.
+            encryption_spec_key_name (Optional[str]):
+                Optional. The Cloud KMS resource identifier of the customer
+                managed encryption key used to protect the job. Has the
+                form:
+                ``projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key``.
+                The key needs to be in the same region as where the compute
+                resource is created.
+
+                If this is set, then all
+                resources created by the BatchPredictionJob will
+                be encrypted with the provided encryption key.
+
+                Overrides encryption_spec_key_name set in aiplatform.init.
             sync (bool):
                 Whether to execute this method synchronously. If False, this method
                 will be executed in concurrent Future and any downstream object will
@@ -459,6 +474,9 @@ def create(
         gapic_batch_prediction_job.output_config = output_config
 
         # Optional Fields
+        gapic_batch_prediction_job.encryption_spec = initializer.global_config.get_encryption_spec(
+            encryption_spec_key_name=encryption_spec_key_name
+        )
 
         if model_parameters:
             gapic_batch_prediction_job.model_parameters = model_parameters